pass kwargs to minimizer

pymc_marketing/mmm/budget_optimizer.py

@@ -23,6 +23,13 @@
 from pymc_marketing.mmm.components.saturation import SaturationTransformation
 
 
+class MinimizeException(Exception):
+    """Custom exception for optimization failure."""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
 class BudgetOptimizer:
     """
     A class for optimizing budget allocation in a marketing mix model.
@@ -112,6 +119,7 @@
         total_budget: float,
         budget_bounds: dict[str, tuple[float, float]] | None = None,
         custom_constraints: dict[Any, Any] | None = None,
+        minimize_kwargs: dict[str, Any] | None = None,
     ) -> tuple[dict[str, float], float]:
         """
         Allocate the budget based on the total budget, budget bounds, and custom constraints.
@@ -136,6 +144,9 @@
             The budget bounds for each channel. Default is None.
         custom_constraints : dict, optional
             Custom constraints for the optimization. Default is None.
+        minimize_kwargs : dict, optional
+            Additional keyword arguments for the `scipy.optimize.minimize` function. If None, default values are used.
+            Method is set to "SLSQP", ftol is set to 1e-9, and maxiter is set to 1_000.
 
         Returns
         -------
@@ -179,19 +190,26 @@
             )
             for channel in self.parameters
         ]
+
+        if minimize_kwargs is None:
+            minimize_kwargs = {
+                "method": "SLSQP",
+                "options": {"ftol": 1e-9, "maxiter": 1_000},
+            }
+
         result = minimize(
             self.objective,
             x0=initial_guess,
             bounds=bounds,
             constraints=constraints,
-            method="SLSQP",
-            options={"ftol": 1e-9, "maxiter": 1000},
+            **minimize_kwargs,
         )
+
         if result.success:
             optimal_budgets = {
                 name: budget
                 for name, budget in zip(self.parameters.keys(), result.x, strict=False)
             }
             return optimal_budgets, -result.fun
         else:
-            raise Exception("Optimization failed: " + result.message)
+            raise MinimizeException(f"Optimization failed: {result.message}")

tests/mmm/test_budget_optimizer.py

@@ -13,14 +13,31 @@
 #   limitations under the License.
 import pytest
 
-from pymc_marketing.mmm.budget_optimizer import BudgetOptimizer
+from pymc_marketing.mmm.budget_optimizer import BudgetOptimizer, MinimizeException
 from pymc_marketing.mmm.components.adstock import _get_adstock_function
 from pymc_marketing.mmm.components.saturation import _get_saturation_function
 
 
 @pytest.mark.parametrize(
-    "total_budget, budget_bounds, parameters, expected_optimal, expected_response",
-    [
+    argnames="total_budget, budget_bounds, parameters, minimize_kwargs, expected_optimal, expected_response",
+    argvalues=[
+        (
+            100,
+            {"channel_1": (0, 50), "channel_2": (0, 50)},
+            {
+                "channel_1": {
+                    "adstock_params": {"alpha": 0.5},
+                    "saturation_params": {"lam": 10, "beta": 0.5},
+                },
+                "channel_2": {
+                    "adstock_params": {"alpha": 0.7},
+                    "saturation_params": {"lam": 20, "beta": 1.0},
+                },
+            },
+            None,
+            {"channel_1": 50.0, "channel_2": 50.0},
+            49.5,
+        ),
         (
             100,
             {"channel_1": (0, 50), "channel_2": (0, 50)},
@@ -34,14 +51,24 @@
                     "saturation_params": {"lam": 20, "beta": 1.0},
                 },
             },
+            {
+                "method": "SLSQP",
+                "options": {"ftol": 1e-8, "maxiter": 1_002},
+            },
             {"channel_1": 50.0, "channel_2": 50.0},
             49.5,
         ),
         # Add more test cases if needed
     ],
+    ids=["default_minimizer_kwargs", "custom_minimizer_kwargs"],
 )
 def test_allocate_budget(
-    total_budget, budget_bounds, parameters, expected_optimal, expected_response
+    total_budget,
+    budget_bounds,
+    parameters,
+    minimize_kwargs,
+    expected_optimal,
+    expected_response,
 ):
     # Initialize Adstock and Saturation Transformations
     adstock = _get_adstock_function(function="geometric", l_max=4)
@@ -52,7 +79,9 @@ def test_allocate_budget(
 
     # Allocate Budget
     optimal_budgets, total_response = optimizer.allocate_budget(
-        total_budget, budget_bounds
+        total_budget=total_budget,
+        budget_bounds=budget_bounds,
+        minimize_kwargs=minimize_kwargs,
     )
 
     # Assert Results
@@ -124,5 +153,5 @@ def test_allocate_budget_infeasible_constraints(
     saturation = _get_saturation_function(function="logistic")
     optimizer = BudgetOptimizer(adstock, saturation, 30, parameters, adstock_first=True)
 
-    with pytest.raises(Exception, match="Optimization failed"):
+    with pytest.raises(MinimizeException, match="Optimization failed"):
         optimizer.allocate_budget(total_budget, budget_bounds, custom_constraints)