Create wrapper to make budget optimizer compatible with multidimensional class (#1652)

* Saving progress

* Changes

* Closing remarks

* Adding plots

Not working yet

* Updating

* Adding quick test

* Solving plot issues

* Adding tests

* Solve issue

Ricky always helping!

Co-Authored-By: Ricardo Vieira <[email protected]>

* Solving test issues

* Modifying example

* Adding more explanations

* Removing comment

* Requested changes

* Replacing cell.

* Other small change.

---------

Co-authored-by: Ricardo Vieira <[email protected]>
Co-authored-by: Will Dean <[email protected]>

Author: 3 people

docs/source/notebooks/mmm/mmm_multidimensional_example.ipynb

@@ -165,7 +165,9 @@ class BudgetOptimizer(BaseModel):
     def __init__(self, **data):
         super().__init__(**data)
         # 1. Prepare model with time dimension for optimization
-        pymc_model = self.mmm_model._set_predictors_for_optimization(self.num_periods)
+        pymc_model = self.mmm_model._set_predictors_for_optimization(
+            self.num_periods
+        )  # TODO: Once multidimensional class becomes the main class.
 
         # 2. Shared variable for total_budget: Use annotation to avoid type checking
         self._total_budget: SharedVariable = shared(
@@ -270,13 +272,20 @@ def _replace_channel_data_by_optimization_variable(self, model: Model) -> Model:
         repeated_budgets_with_carry_over_shape.insert(
             date_dim_idx, num_periods + max_lag
         )
+
+        # Get the dtype from the model's channel_data to ensure type compatibility
+        channel_data_dtype = model["channel_data"].dtype
+
         repeated_budgets_with_carry_over = pt.zeros(
-            repeated_budgets_with_carry_over_shape
+            repeated_budgets_with_carry_over_shape,
+            dtype=channel_data_dtype,  # Use the same dtype as channel_data
         )
         set_idxs = (*((slice(None),) * date_dim_idx), slice(None, num_periods))
         repeated_budgets_with_carry_over = repeated_budgets_with_carry_over[
             set_idxs
-        ].set(repeated_budgets)
+        ].set(
+            pt.cast(repeated_budgets, channel_data_dtype)
+        )  # Cast to ensure type compatibility
         repeated_budgets_with_carry_over.name = "repeated_budgets_with_carry_over"
 
         # Freeze dims & data in the underlying PyMC model

pymc_marketing/mmm/budget_optimizer.py

@@ -17,6 +17,7 @@
 
 import json
 import warnings
+from collections.abc import Sequence
 from copy import deepcopy
 from typing import Any, Literal
 
@@ -29,9 +30,11 @@
 import xarray as xr
 from pymc.model.fgraph import clone_model as cm
 from pymc.util import RandomState
+from scipy.optimize import OptimizeResult
 
 from pymc_marketing.mmm import SoftPlusHSGP
 from pymc_marketing.mmm.additive_effect import MuEffect, create_event_mu_effect
+from pymc_marketing.mmm.budget_optimizer import OptimizerCompatibleModelWrapper
 from pymc_marketing.mmm.components.adstock import (
     AdstockTransformation,
     adstock_from_dict,
@@ -45,6 +48,11 @@
 from pymc_marketing.mmm.plot import MMMPlotSuite
 from pymc_marketing.mmm.scaling import Scaling, VariableScaling
 from pymc_marketing.mmm.tvp import infer_time_index
+from pymc_marketing.mmm.utility import UtilityFunctionType, average_response
+from pymc_marketing.mmm.utils import (
+    add_noise_to_channel_allocation,
+    create_zero_dataset,
+)
 from pymc_marketing.model_builder import ModelBuilder, _handle_deprecate_pred_argument
 from pymc_marketing.model_config import parse_model_config
 from pymc_marketing.model_graph import deterministics_to_flat
@@ -945,14 +953,15 @@ def build_model(
             channel_data_.name = "channel_data_scaled"
             channel_data_.dims = ("date", *self.dims, "channel")
 
-            ## Hot fix for target data meanwhile pymc allows for internal scaling `https://github.com/pymc-devs/pymc/pull/7656`
             target_dim_handler = create_dim_handler(("date", *self.dims))
-            target_data_scaled = pm.Deterministic(
-                name="target_scaled",
-                var=_target
-                / target_dim_handler(_target_scale, self.scalers._target.dims),
-                dims=("date", *self.dims),
+
+            target_data_scaled = _target / target_dim_handler(
+                _target_scale, self.scalers._target.dims
             )
+            target_data_scaled.name = "target_scaled"
+            target_data_scaled.dims = ("date", *self.dims)
+            ## TODO: Find a better way to save it or access it in the pytensor graph.
+            self.target_data_scaled = target_data_scaled
 
             for mu_effect in self.mu_effects:
                 mu_effect.create_data(self)
@@ -1417,3 +1426,125 @@ def create_sample_kwargs(
     # Update with additional keyword arguments
     sampler_config.update(kwargs)
     return sampler_config
+
+
+class MultiDimensionalBudgetOptimizerWrapper(OptimizerCompatibleModelWrapper):
+    """Wrapper for the BudgetOptimizer to handle multi-dimensional model."""
+
+    def __init__(self, model: MMM, start_date: str, end_date: str):
+        self.model_class = model
+        self.start_date = start_date
+        self.end_date = end_date
+        # Compute the number of periods to allocate budget for
+        self.zero_data = create_zero_dataset(
+            model=self.model_class, start_date=start_date, end_date=end_date
+        )
+        self.num_periods = len(self.zero_data[self.model_class.date_column].unique())
+        # Adding missing dependencies for compatibility with BudgetOptimizer
+        self._channel_scales = 1.0
+
+    def __getattr__(self, name):
+        """Delegate attribute access to the wrapped MMM model."""
+        try:
+            # First, try to get the attribute from the wrapper itself
+            return object.__getattribute__(self, name)
+        except AttributeError:
+            # If not found, delegate to the wrapped model
+            try:
+                return getattr(self.model_class, name)
+            except AttributeError as e:
+                # Raise an AttributeError if the attribute is not found in either
+                raise AttributeError(
+                    f"'{type(self).__name__}' object and its wrapped 'MMM' object have no attribute '{name}'"
+                ) from e
+
+    def _set_predictors_for_optimization(self, num_periods: int) -> pm.Model:
+        """Return the respective PyMC model with any predictors set for optimization."""
+        # Use the model's method for transformation
+        dataset_xarray = self._posterior_predictive_data_transformation(
+            X=self.zero_data,
+            include_last_observations=False,
+        )
+
+        # Use the model's method to set data
+        pymc_model = self._set_xarray_data(
+            dataset_xarray=dataset_xarray,
+            clone_model=True,  # Ensure we work on a clone
+        )
+
+        # Use the model's mu_effects and set data using the model instance
+        for mu_effect in self.mu_effects:
+            mu_effect.set_data(self, pymc_model, dataset_xarray)
+
+        return pymc_model
+
+    def optimize_budget(
+        self,
+        budget: float | int,
+        budget_bounds: xr.DataArray | dict[str, tuple[float, float]] | None = None,
+        response_variable: str = "total_media_contribution_original_scale",
+        utility_function: UtilityFunctionType = average_response,
+        constraints: Sequence[dict[str, Any]] = (),
+        default_constraints: bool = True,
+        **minimize_kwargs,
+    ) -> tuple[xr.DataArray, OptimizeResult]:
+        """Optimize the budget allocation for the model."""
+        from pymc_marketing.mmm.budget_optimizer import BudgetOptimizer
+
+        allocator = BudgetOptimizer(
+            num_periods=self.num_periods,
+            utility_function=utility_function,
+            response_variable=response_variable,
+            custom_constraints=constraints,
+            default_constraints=default_constraints,
+            model=self,  # Pass the wrapper instance itself to the BudgetOptimizer
+        )
+
+        return allocator.allocate_budget(
+            total_budget=budget,
+            budget_bounds=budget_bounds,
+            **minimize_kwargs,
+        )
+
+    def sample_response_distribution(
+        self,
+        allocation_strategy: xr.DataArray,
+        noise_level: float = 0.001,
+    ) -> az.InferenceData:
+        """Generate synthetic dataset and sample posterior predictive based on allocation.
+
+        Parameters
+        ----------
+        allocation_strategy : DataArray
+            The allocation strategy for the channels.
+        noise_level : float
+            The relative level of noise to add to the data allocation.
+
+        Returns
+        -------
+        az.InferenceData
+            The posterior predictive samples based on the synthetic dataset.
+        """
+        data = create_zero_dataset(
+            model=self,
+            start_date=self.start_date,
+            end_date=self.end_date,
+            channel_xr=allocation_strategy.to_dataset(dim="channel"),
+        )
+
+        data_with_noise = add_noise_to_channel_allocation(
+            df=data,
+            channels=self.channel_columns,
+            rel_std=noise_level,
+            seed=42,
+        )
+
+        constant_data = allocation_strategy.to_dataset(name="allocation")
+
+        return self.sample_posterior_predictive(
+            X=data_with_noise,
+            extend_idata=False,
+            include_last_observations=True,
+            var_names=["y", "channel_contribution_original_scale"],
+            progressbar=False,
+        ).merge(constant_data)