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

[cluhmann]

Description

BudgetOptimizer.allocate_budget() now accepts x0 as an argument.

Related Issue

• Closes BudgetOptimizer.allocate_budget() doc string doesn't specify default value of x0 #1564
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass. Feel free to comment pre-commit.ci autofix to auto-fix.
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks) using numpydoc format.
• If you are a pro: each commit corresponds to a relevant logical change

x0 is now accepted via the minimize_kwargs (as the docstring implies is possible). I looked into adding a corresponding test, but these lines in tests/mmm/test_budget_optimizer::test_allocate_budget_custom_minimize_args() confuse me. Any guidance?

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)]
assert kwargs["method"] == minimize_kwargs["method"]
assert kwargs["options"] == minimize_kwargs["options"]

---

📚 Documentation preview 📚: https://pymc-marketing--1565.org.readthedocs.build/en/1565/

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 93.51%. Comparing base (8475c1c) to head (c8a28ec).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1565   +/-   ##
=======================================
  Coverage   93.51%   93.51%           
=======================================
  Files          55       55           
  Lines        6341     6342    +1     
=======================================
+ Hits         5930     5931    +1     
  Misses        411      411           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ricardoV94]

I don't think x0 should be hidden in minimize kwargs, I would make it an explicit argument of allocate_budget

[ricardoV94]

I would check if it has the expected shape (budgets_size,), as the compiled functions have trust_input=True and won't check if x0 has the promised shape

[cetagostini]

very good point here!

[cetagostini]

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)]
assert kwargs["method"] == minimize_kwargs["method"]
assert kwargs["options"] == minimize_kwargs["options"]

Just a hardcoded check based on the data, we could change this to be more dynamic and check if values are equal based on the data, not hardcoded 50/50 🙃 Does that makes it clear?

[ricardoV94]

Nit: no real need to copy

Suggested change

	minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS.copy()
	minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS

[ricardoV94]

I think it should be a np.ndarray not a DataArray

[ricardoV94]

Should also validate dtype

[cluhmann]

Just need to write some tests.

[ricardoV94]

Specifically the x0.dtype must match self._budgets_flat.type.dtype. I think there's a shortcut you can use for the shape and dtype validation: x0 = self._budgets_flat.type.filter(x0), which will do any allowed casting or raise if things cannot be safely cast (and also checks shape)