From da963d2c357ff28e70fec940d7cf5cc7e5ac91cd Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Wed, 12 Jun 2024 08:20:36 +0200 Subject: [PATCH] pass kwargs to minimizer (#737) --- pymc_marketing/mmm/budget_optimizer.py | 30 +++++++-- tests/mmm/test_budget_optimizer.py | 86 ++++++++++++++++++++++++-- 2 files changed, 104 insertions(+), 12 deletions(-) diff --git a/pymc_marketing/mmm/budget_optimizer.py b/pymc_marketing/mmm/budget_optimizer.py index 73cf66cc..8fa4bd16 100644 --- a/pymc_marketing/mmm/budget_optimizer.py +++ b/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,14 +190,21 @@ 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 @@ -194,4 +212,4 @@ def allocate_budget( } return optimal_budgets, -result.fun else: - raise Exception("Optimization failed: " + result.message) + raise MinimizeException(f"Optimization failed: {result.message}") diff --git a/tests/mmm/test_budget_optimizer.py b/tests/mmm/test_budget_optimizer.py index 637de20a..97ff6e70 100644 --- a/tests/mmm/test_budget_optimizer.py +++ b/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)