Permit BudgetOptimizer.allocate_budget() to take x0 as an argument (#1565)

* formatting

* expose x0 as an arg in allocate_budget

* change typehint and validate dtype

* streamline shape and dtype validation

* add new bo test with manually set x0

---------

Co-authored-by: Carlos Trujillo <59846724+cetagostini@users.noreply.github.com>

Author: cluhmann

pymc_marketing/mmm/budget_optimizer.py

@@ -373,6 +373,7 @@ def allocate_budget(
         self,
         total_budget: float,
         budget_bounds: DataArray | dict[str, tuple[float, float]] | None = None,
+        x0: np.ndarray | None = None,
         minimize_kwargs: dict[str, Any] | None = None,
         return_if_fail: bool = False,
     ) -> tuple[DataArray, OptimizeResult]:
@@ -391,8 +392,11 @@ def allocate_budget(
             - If None, default bounds of [0, total_budget] per channel are assumed.
             - If a dict, must map each channel to (low, high) budget pairs (only valid if there's one dimension).
             - If an xarray.DataArray, must have dims (*budget_dims, "bound"), specifying [low, high] per channel cell.
+        x0 : np.ndarray, optional
+            Initial guess. Array of real elements of size (n,), where n is the number of driver budgets to optimize. If
+            None, the total budget is spread uniformly across all drivers to be optimized.
         minimize_kwargs : dict, optional
-            Extra kwargs for `scipy.optimize.minimize`. Defaults to method "SLSQP",
+            Extra kwargs for `scipy.optimize.minimize`. Defaults to method="SLSQP",
             ftol=1e-9, maxiter=1_000.
         return_if_fail : bool, optional
             Return output even if optimization fails. Default is False.
@@ -409,8 +413,16 @@ def allocate_budget(
         MinimizeException
             If the optimization fails for any reason, the exception message will contain the details.
         """
+        # set total budget
         self._total_budget.set_value(np.asarray(total_budget, dtype="float64"))
 
+        # coordinate user-provided and default minimize_kwargs
+        if minimize_kwargs is None:
+            minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS
+        else:
+            # Merge with defaults (preferring user-supplied keys)
+            minimize_kwargs = {**self.DEFAULT_MINIMIZE_KWARGS, **minimize_kwargs}
+
         # 1. Process budget bounds
         if budget_bounds is None:
             warnings.warn(
@@ -466,21 +478,21 @@ def allocate_budget(
         else:
             budgets_size = self.budgets_to_optimize.sum().item()
 
-        # 5. Create an initial guess
-        initial_guess = np.ones(budgets_size) * (total_budget / budgets_size)
-        initial_guess = initial_guess.astype(self._budgets_flat.type.dtype)
+        # 5. Construct the initial guess (x0) if not provided
+        if x0 is None:
+            x0 = np.ones(budgets_size) * (total_budget / budgets_size).astype(
+                self._budgets_flat.type.dtype
+            )
 
-        if minimize_kwargs is None:
-            minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS.copy()
-        else:
-            # Merge with defaults (preferring user-supplied keys)
-            minimize_kwargs = {**self.DEFAULT_MINIMIZE_KWARGS, **minimize_kwargs}
+        # filter x0 based on shape/type of self._budgets_flat
+        # will raise a TypeError if x0 does not have acceptable shape and/or type
+        x0 = self._budgets_flat.type.filter(x0)
 
         # 6. Run the SciPy optimizer
         result = minimize(
             fun=self._compiled_functions[self.utility_function]["objective_and_grad"],
+            x0=x0,
             jac=True,
-            x0=initial_guess,
             bounds=bounds,
             constraints=self._compiled_constraints,
             **minimize_kwargs,

tests/mmm/test_budget_optimizer.py

@@ -59,11 +59,36 @@ def dummy_df():
 
 
 @pytest.mark.parametrize(
-    argnames="total_budget, budget_bounds, parameters, minimize_kwargs, expected_optimal, expected_response",
+    argnames="total_budget, budget_bounds, x0, parameters, minimize_kwargs, expected_optimal, expected_response",
     argvalues=[
         (
             100,
             None,
+            None,
+            {
+                "saturation_params": {
+                    "lam": np.array(
+                        [[[0.1, 0.2], [0.3, 0.4]], [[0.5, 0.6], [0.7, 0.8]]]
+                    ),  # dims: chain, draw, channel
+                    "beta": np.array(
+                        [[[0.5, 1.0], [0.5, 1.0]], [[0.5, 1.0], [0.5, 1.0]]]
+                    ),  # dims: chain, draw, channel
+                },
+                "adstock_params": {
+                    "alpha": np.array(
+                        [[[0.5, 0.7], [0.5, 0.7]], [[0.5, 0.7], [0.5, 0.7]]]
+                    )  # dims: chain, draw, channel
+                },
+            },
+            None,
+            {"channel_1": 54.78357587906867, "channel_2": 45.21642412093133},
+            48.8,
+        ),
+        # set x0 manually
+        (
+            100,
+            None,
+            np.array([50, 50]),
             {
                 "saturation_params": {
                     "lam": np.array(
@@ -91,6 +116,7 @@ def dummy_df():
                 channel=["channel_1", "channel_2"],
                 bound=["lower", "upper"],
             ),
+            None,
             {
                 "saturation_params": {
                     "lam": np.array(
@@ -121,6 +147,7 @@ def dummy_df():
                 channel=["channel_1", "channel_2"],
                 bound=["lower", "upper"],
             ),
+            None,
             {
                 "saturation_params": {
                     "lam": np.array(
@@ -142,11 +169,17 @@ def dummy_df():
             2.38e-10,
         ),
     ],
-    ids=["default_minimizer_kwargs", "custom_minimizer_kwargs", "zero_total_budget"],
+    ids=[
+        "default_minimizer_kwargs",
+        "manually_set_x0",
+        "custom_minimizer_kwargs",
+        "zero_total_budget",
+    ],
 )
 def test_allocate_budget(
     total_budget,
     budget_bounds,
+    x0,
     parameters,
     minimize_kwargs,
     expected_optimal,
@@ -184,6 +217,7 @@ def test_allocate_budget(
     optimal_budgets, optimization_res = optimizer.allocate_budget(
         total_budget=total_budget,
         budget_bounds=budget_bounds,
+        x0=x0,
         minimize_kwargs=minimize_kwargs,
     )