Handle optimization subset and automatically set predictors

pymc_marketing/mmm/budget_optimizer.py

@@ -15,20 +15,20 @@
 """Budget optimization module."""
 
 import warnings
-from typing import Any, ClassVar, Sequence
+from collections.abc import Sequence
+from typing import Any, ClassVar
 
 import numpy as np
 import pytensor.tensor as pt
 from pydantic import BaseModel, ConfigDict, Field
+from pymc import inputvars
 from pymc.logprob.utils import rvs_in_graph
 from pymc.model.transform.optimization import freeze_dims_and_data
 from pytensor import clone_replace, function
 from pytensor.graph import rewrite_graph, vectorize_graph
 from scipy.optimize import minimize
+from xarray import DataArray
 
-from pymc_marketing.mmm.mmm import MMM
-from pymc_marketing.mmm.components.adstock import AdstockTransformation
-from pymc_marketing.mmm.components.saturation import SaturationTransformation
 from pymc_marketing.mmm.constraints import (
     build_constraint,
     build_default_sum_constraint,
@@ -71,7 +71,7 @@ class BudgetOptimizer(BaseModel):
         description="The number of time units at time granularity which the budget is to be allocated.",
     )
 
-    hmm_model: MMM = Field(
+    hmm_model: Any = Field(
         ...,
         description="The marketing mix model to optimize.",
         arbitrary_types_allowed=True,
@@ -97,6 +97,11 @@ class BudgetOptimizer(BaseModel):
         default=True,
         description="Whether to set the default sum constraint to the optimizer.",
     )
+    opt_mask: DataArray | None = Field(
+        ...,
+        description="Mask defining a subset of budgets that should be optimized. "
+        "Non optimized budgets are fixed to 0.",
+    )
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -107,14 +112,43 @@ class BudgetOptimizer(BaseModel):
 
     def __init__(self, **data):
         super().__init__(**data)
-        self._num_channels = len(self.hmm_model.model.coords["channel"])
+        self._pymc_model = self.hmm_model._set_predictors_for_optimization(
+            self.num_periods
+        )
+        self._coords = self._pymc_model._coords
         self._compiled_functions = {}
         self._compile_objective_and_grad()
         self._constraints = {}
-        self.set_constraints(default=self.default_constraints, constraints=self.custom_constraints)
+        self.set_constraints(
+            default=self.default_constraints, constraints=self.custom_constraints
+        )
+
+    def _create_budget_variable(self):
+        model = self._pymc_model
+        coords = self._coords
+
+        budgets_shape = [
+            len(coords[dim])
+            for dim in model.named_vars_to_dims["channel_data"]
+            if dim != "date"
+        ]
+        if self.opt_mask is not None:
+            size_budgets = self.opt_mask.sum().item()
+        else:
+            size_budgets = np.prod(budgets_shape)
+        budgets_flat = pt.tensor("budgets_flat", shape=(size_budgets,))
+
+        if self.opt_mask is not None:
+            budgets = pt.zeros(budgets_shape)[
+                np.asarray(self.opt_mask).astype(bool)
+            ].set(budgets_flat)
+        else:
+            budgets = budgets_flat.reshape(budgets_shape)
+
+        return budgets
 
     def set_constraints(self, constraints, default=None):
-        """ set constraints """
+        """Set constraints"""
         self._constraints = {}
         if default is None:
             default = False if constraints else True
@@ -126,28 +160,31 @@ def set_constraints(self, constraints, default=None):
                 constraint_type=c.get("constraint_type", "eq"),
             )
             self._constraints[c["key"]] = new_constraint
-        
+
         if default:
             self._constraints["default"] = build_default_sum_constraint("default")
 
-        self._compiled_constraints = compile_constraints_for_scipy(constraints=self._constraints, optimizer=self)
+        budgets = self._create_budget_variable()
+        self._compiled_constraints = compile_constraints_for_scipy(
+            budgets=budgets, constraints=self._constraints, optimizer=self
+        )
 
     def _compile_objective_and_grad(self):
         """Compile the objective function and its gradient using symbolic computation."""
-        budgets_sym = pt.vector("budgets", shape=(self._num_channels,))
-
-        response_distribution = self.extract_response_distribution(budgets=budgets_sym)
+        budgets = self._create_budget_variable()
+        response_distribution = self.extract_response_distribution(budgets=budgets)
 
         objective_value = -self.utility_function(
-            samples=response_distribution, budgets=budgets_sym
+            samples=response_distribution, budgets=budgets
         )
 
         # Compute gradient symbolically
-        grad_obj = pt.grad(objective_value, budgets_sym)
+        [budgets_flat] = inputvars([budgets])
+        grad_obj = pt.grad(objective_value, budgets_flat)
 
         # Compile the functions
-        utility_func = function([budgets_sym], objective_value)
-        grad_func = function([budgets_sym], grad_obj)
+        utility_func = function([budgets_flat], objective_value)
+        grad_func = function([budgets_flat], grad_obj)
 
         # Cache the compiled functions
         self._compiled_functions[self.utility_function] = {
@@ -169,27 +206,38 @@ def extract_response_distribution(
         self, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
         """Extract the response graph, conditioned on the posterior draws and a placeholder budget variable."""
-        if not (isinstance(budgets, pt.TensorVariable) and budgets.type.ndim == 1):
-            raise ValueError("budgets must be a 1D TensorVariable")
+        if not (isinstance(budgets, pt.TensorVariable)):  # and budgets.type.ndim == 1):
+            raise ValueError("budgets must be a TensorVariable")
 
-        model = self.hmm_model.model
+        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
 
         # Freeze all but channel dims for a more succinct graph
-        # model = freeze_dims_and_data(
-        #     model, data=[], dims=[dim for dim in model.coords if dim != "date"]
-        # )
+        model = freeze_dims_and_data(
+            model, data=[], dims=[dim for dim in self._coords if dim != "date"]
+        )
 
         response_variable = model[self.response_variable]
 
         # Replicate the budget over num_periods and append zeros to also quantify carry-over effects
-        n_channels = len(model.coords["channel"])
-        budgets_tiled = pt.broadcast_to(budgets, (num_periods, n_channels))
-        budgets_full = pt.zeros((num_periods + max_lag, n_channels))
+        channel_data_dims = model.named_vars_to_dims["channel_data"]
+        date_dim_idx = list(channel_data_dims).index("date")
+
+        budgets_tiled_shape = list(tuple(budgets.shape))
+        budgets_tiled_shape.insert(date_dim_idx, num_periods)
+        budgets_tiled = pt.broadcast_to(
+            pt.expand_dims(budgets, date_dim_idx), budgets_tiled_shape
+        )
+        # print(f"{budgets_tiled.type=}, {budgets_tiled_shape=}, {channel_data_dims=}")
+
+        budget_full_shape = list(tuple(budgets.shape))
+        budget_full_shape.insert(date_dim_idx, num_periods + max_lag)
+        budgets_full = pt.zeros(budget_full_shape)
         # budgets_full = budgets_full[:num_periods, :].set(budgets_tiled)
-        budgets_full = pt.set_subtensor(budgets_full[:num_periods, :], budgets_tiled)
+        set_idxs = (*((slice(None),) * date_dim_idx), slice(None, num_periods))
+        budgets_full = pt.set_subtensor(budgets_full[set_idxs], budgets_tiled)
         budgets_full.name = "budgets_full"
 
         # Replace model free_RVs by placeholder variables
@@ -306,7 +354,9 @@ def allocate_budget(
                 }
             )
 
-        initial_guess = np.ones(self._num_channels) * total_budget / self._num_channels
+        [budgets_size] = inputvars(self._create_budget_variable())[0].type.shape
+
+        initial_guess = np.ones(budgets_size) * total_budget / budgets_size
         bounds = [
             (
                 (budget_bounds[channel][0], budget_bounds[channel][1])
@@ -325,7 +375,7 @@ def allocate_budget(
         result = minimize(
             fun=self._objective,
             x0=initial_guess,
-            bounds=bounds,
+            # bounds=bounds,
             constraints=constraints_for_scipy,
             jac=self._gradient,
             **minimize_kwargs,

pymc_marketing/mmm/constraints.py

@@ -1,8 +1,22 @@
+#   Copyright 2025 The PyMC Labs Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
 # constraints.py
 
 import pytensor.tensor as pt
+from pymc.pytensorf import inputvars, rewrite_pregrad
 from pytensor import function
-from pymc.pytensorf import rewrite_pregrad
+
 
 def auto_jacobian(
     constraint_fun,
@@ -11,11 +25,15 @@ def auto_jacobian(
     Given a symbolic constraint function constraint_fun(budgets_sym, total_budget_sym, optimizer),
     return a symbolic jacobian function that depends on the same variables.
     """
+
     def _jac(budgets_sym, total_budget_sym, optimizer):
         _fun = constraint_fun(budgets_sym, total_budget_sym, optimizer)
-        return pt.grad(rewrite_pregrad(_fun), budgets_sym)
+        [budgets_flat] = inputvars([budgets_sym])
+        return pt.grad(rewrite_pregrad(_fun), budgets_flat)
+
     return _jac
 
+
 def build_constraint(
     key: str,
     constraint_type: str,
@@ -43,10 +61,12 @@ def build_constraint(
         "sym_jac": constraint_jac,
     }
 
+
 def build_default_sum_constraint(key: str = "default"):
     """
-    Returns a constraint dict that enforces sum(budgets) == total_budget.
+    Return a constraint dict that enforces sum(budgets) == total_budget.
     """
+
     def _constraint_fun(budgets_sym, total_budget_sym, optimizer):
         return pt.sum(budgets_sym) - total_budget_sym
 
@@ -57,12 +77,13 @@ def _constraint_fun(budgets_sym, total_budget_sym, optimizer):
         constraint_jac=None,
     )
 
-def compile_constraints_for_scipy(constraints, optimizer):
-    """ compile constraints for scipy """
+
+def compile_constraints_for_scipy(budgets, constraints, optimizer):
+    """Compile constraints for scipy."""
     compiled_constraints = []
 
-    budgets_sym = pt.vector("budgets")
-    total_budget_sym = pt.scalar("total_budget")
+    [budgets_flat] = inputvars([budgets])
+    total_budget = pt.scalar("total_budget")
 
     for c in constraints.values() if isinstance(constraints, dict) else constraints:
         ctype = c["type"]
@@ -71,19 +92,21 @@ def compile_constraints_for_scipy(constraints, optimizer):
 
         # Compile symbolic => python callables
         compiled_fun = function(
-            inputs=[budgets_sym, total_budget_sym],
-            outputs=sym_fun(budgets_sym, total_budget_sym, optimizer),
+            inputs=[budgets_flat, total_budget],
+            outputs=sym_fun(budgets, total_budget, optimizer),
             on_unused_input="ignore",
         )
         compiled_jac = function(
-            inputs=[budgets_sym, total_budget_sym],
-            outputs=sym_jac(budgets_sym, total_budget_sym, optimizer),
+            inputs=[budgets_flat, total_budget],
+            outputs=sym_jac(budgets, total_budget, optimizer),
             on_unused_input="ignore",
         )
 
-        compiled_constraints.append({
-            "type": ctype,
-            "fun": compiled_fun,
-            "jac": compiled_jac,
-        })
+        compiled_constraints.append(
+            {
+                "type": ctype,
+                "fun": compiled_fun,
+                "jac": compiled_jac,
+            }
+        )
     return compiled_constraints

pymc_marketing/mmm/mmm.py

@@ -16,7 +16,8 @@
 import json
 import logging
 import warnings
-from typing import Annotated, Any, Literal, Sequence
+from collections.abc import Sequence
+from typing import Annotated, Any, Literal
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -32,7 +33,6 @@
 
 from pymc_marketing.hsgp_kwargs import HSGPKwargs
 from pymc_marketing.mmm.base import BaseValidateMMM
-
 from pymc_marketing.mmm.components.adstock import (
     AdstockTransformation,
     adstock_from_dict,
@@ -2232,6 +2232,22 @@ def sample_response_distribution(
             progressbar=False,
         ).merge(constant_data)
 
+    def _set_predictors_for_optimization(self, num_periods: int) -> pm.Model:
+        """Return the respective PyMC model with any predictors set for optimization."""
+        model = self.model
+        # Models with HSGP have a time_index data variable
+        if "time_index" in model.named_vars:
+            model = model.copy()
+            start_date = model["time_index"].get_value(borrow=True)[-1]
+            training_dates = (
+                np.arange(num_periods + self.adstock.l_max) + start_date + 1
+            )
+            # Consider using pm.set_data, but then we need new coordinates
+            model["time_index"].set_value(
+                training_dates.astype(model["time_index"].type.dtype)
+            )
+        return model
+
     def optimize_budget(
         self,
         budget: float | int,
@@ -2293,7 +2309,6 @@ def optimize_budget(
         ValueError
             If the noise level is not a float.
         """
-
         from pymc_marketing.mmm.budget_optimizer import BudgetOptimizer
 
         allocator = BudgetOptimizer(