remove spend scaling step from budget optimizer (pymc-labs#1070)

* remove spend scaling step * re-run nb
jsnyde0 · Sep 30, 2024 · 210a53b · 210a53b
1 parent 0592fcb
commit 210a53b
Show file tree

Hide file tree

Showing 2 changed files with 119 additions and 109 deletions.
diff --git a/docs/source/notebooks/mmm/mmm_budget_allocation_example.ipynb b/docs/source/notebooks/mmm/mmm_budget_allocation_example.ipynb
diff --git a/pymc_marketing/mmm/mmm.py b/pymc_marketing/mmm/mmm.py
@@ -2239,28 +2239,14 @@ def plot_budget_allocation(
             The matplotlib figure object and axis containing the plot.
 
         """
-        if original_scale:
-            channel_contributions = (
-                samples["channel_contributions"]
-                .mean(dim=["sample"])
-                .mean(dim=["date"])
-                .values
-                * self.get_target_transformer()["scaler"].scale_
-            )
+        channel_contributions = (
+            samples["channel_contributions"].mean(dim=["date", "sample"]).to_numpy()
+        )
 
-            allocate_spend = (
-                np.array(list(self.optimal_allocation_dict.values()))
-                * self.channel_transformer["scaler"].scale_
-            )
+        if original_scale:
+            channel_contributions *= self.get_target_transformer()["scaler"].scale_
 
-        else:
-            channel_contributions = (
-                samples["channel_contributions"]
-                .mean(dim=["sample"])
-                .mean(dim=["date"])
-                .values
-            )
-            allocate_spend = np.array(list(self.optimal_allocation_dict.values()))
+        allocated_spend = np.array(list(self.optimal_allocation_dict.values()))
 
         if ax is None:
             fig, ax = plt.subplots(figsize=figsize)
@@ -2274,11 +2260,11 @@ def plot_budget_allocation(
 
         bars1 = ax.bar(
             index,
-            allocate_spend,
+            allocated_spend,
             bar_width,
-            color="b",
+            color="C0",
             alpha=opacity,
-            label="Allocate Spend",
+            label="Allocated Spend",
         )
 
         ax2 = ax.twinx()
@@ -2287,19 +2273,19 @@ def plot_budget_allocation(
             index + bar_width,
             channel_contributions,
             bar_width,
-            color="r",
+            color="C1",
             alpha=opacity,
             label="Channel Contributions",
         )
 
         ax.set_xlabel("Channels")
-        ax.set_ylabel("Allocate Spend", color="b")
+        ax.set_ylabel("Allocate Spend", color="C0")
         ax.tick_params(axis="x", rotation=90)
         ax.set_xticks(index + bar_width / 2)
         ax.set_xticklabels(self.channel_columns)
 
-        ax.set_ylabel("Allocate Spend", color="b", labelpad=10)
-        ax2.set_ylabel("Channel Contributions", color="r", labelpad=10)
+        ax.set_ylabel("Allocate Spend", color="C0", labelpad=10)
+        ax2.set_ylabel("Channel Contributions", color="C1", labelpad=10)
 
         ax.grid(False)
         ax2.grid(False)