add new bo test with manually set x0

Author: cluhmann

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,
     )