pass kwargs to minimizer (#737)

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 @@ def allocate_budget(
         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 @@ def allocate_budget(
             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
         -------
@@ -160,7 +171,7 @@ def allocate_budget(
         if custom_constraints is None:
             constraints = {"type": "eq", "fun": lambda x: np.sum(x) - total_budget}
             warnings.warn(
-                "Using default equaliy constraint: The sum of all budgets should be equal to the total budget.",
+                "Using default equality constraint: The sum of all budgets should be equal to the total budget.",
                 stacklevel=2,
             )
         else:
@@ -170,7 +181,7 @@ def allocate_budget(
                 constraints = custom_constraints
 
         num_channels = len(self.parameters.keys())
-        initial_guess = [total_budget // num_channels] * num_channels
+        initial_guess = np.ones(num_channels) * total_budget / num_channels
         bounds = [
             (
                 (budget_bounds[channel][0], budget_bounds[channel][1])
@@ -179,19 +190,26 @@ def allocate_budget(
             )
             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,
+            fun=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

@@ -11,16 +11,19 @@
 #   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.
+from unittest.mock import patch
+
+import numpy as np
 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)},
@@ -34,14 +37,41 @@
                     "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)},
+            {
+                "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},
+                },
+            },
+            {
+                "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 +82,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
@@ -94,6 +126,48 @@ def test_allocate_budget_zero_total(
     assert total_response == pytest.approx(expected_response, abs=1e-1)
 
 
+@patch("pymc_marketing.mmm.budget_optimizer.minimize")
+def test_allocate_budget_custom_minimize_args(minimize_mock) -> None:
+    total_budget = 100
+    budget_bounds = {"channel_1": (0.0, 50.0), "channel_2": (0.0, 50.0)}
+    parameters = {
+        "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},
+        },
+    }
+    minimize_kwargs = {
+        "method": "SLSQP",
+        "options": {"ftol": 1e-8, "maxiter": 1_002},
+    }
+
+    adstock = _get_adstock_function(function="geometric", l_max=4)
+    saturation = _get_saturation_function(function="logistic")
+    optimizer = BudgetOptimizer(adstock, saturation, 30, parameters, adstock_first=True)
+    optimizer.allocate_budget(
+        total_budget, budget_bounds, minimize_kwargs=minimize_kwargs
+    )
+
+    kwargs = minimize_mock.call_args_list[0].kwargs
+
+    np.testing.assert_array_equal(x=kwargs["x0"], y=np.array([50.0, 50.0]))
+    assert kwargs["bounds"] == [(0.0, 50.0), (0.0, 50.0)]
+    # default constraint constraints = {"type": "eq", "fun": lambda x: np.sum(x) - total_budget}
+    assert kwargs["constraints"]["type"] == "eq"
+    assert (
+        kwargs["constraints"]["fun"](np.array([total_budget / 2, total_budget / 2]))
+        == 0.0
+    )
+    assert kwargs["constraints"]["fun"](np.array([100.0, 0.0])) == 0.0
+    assert kwargs["constraints"]["fun"](np.array([0.0, 0.0])) == -total_budget
+    assert kwargs["method"] == minimize_kwargs["method"]
+    assert kwargs["options"] == minimize_kwargs["options"]
+
+
 @pytest.mark.parametrize(
     "total_budget, budget_bounds, parameters, custom_constraints",
     [
@@ -124,5 +198,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)