Scale input

pymc_marketing/mmm/budget_optimizer.py

@@ -144,7 +144,6 @@ def _create_budget_variable(self):
             ].set(budgets_flat)
         else:
             budgets = budgets_flat.reshape(budgets_shape)
-
         return budgets
 
     def set_constraints(self, constraints, default=None):
@@ -209,10 +208,13 @@ def extract_response_distribution(
         if not (isinstance(budgets, pt.TensorVariable)):  # and budgets.type.ndim == 1):
             raise ValueError("budgets must be a TensorVariable")
 
+        num_periods = self.num_periods
         model = self._pymc_model
         posterior = self.hmm_model.idata.posterior  # type: ignore
         max_lag = self.hmm_model.adstock.l_max
-        num_periods = self.num_periods
+        channel_scales = self.hmm_model.scaler.input_scales[
+            self.hmm_model.channel_columns
+        ].values
 
         # Freeze all but channel dims for a more succinct graph
         model = freeze_dims_and_data(
@@ -227,6 +229,8 @@ def extract_response_distribution(
 
         budgets_tiled_shape = list(tuple(budgets.shape))
         budgets_tiled_shape.insert(date_dim_idx, num_periods)
+        # TODO: If scales become part of the model, we don't need to transform it here
+        budgets /= channel_scales[:, None]
         budgets_tiled = pt.broadcast_to(
             pt.expand_dims(budgets, date_dim_idx), budgets_tiled_shape
         )