Fix BudgetOptimizer docstrings

[TeemuSailynoja]

Description

Related Issue

• Closes Wrong BudgetOptimizer docstring #1661

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

---

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

[TeemuSailynoja]

This perhaps would need an example to specify the required shape for the mask.

[cetagostini-wise]

Take a look here:

import numpy as np
import pandas as pd
import xarray as xr

# Create simulated data with only date and channel dimensions
channels = ['TV', 'Radio', 'Social', 'Display', 'Search']  # All available channels
dates = pd.date_range('2023-01-01', periods=10, freq='W')

# Create sample channel_data (shape: date, channel)
channel_data = xr.DataArray(
    np.random.rand(len(dates), len(channels)),
    coords={
        'date': dates,
        'channel': channels  
    },
    dims=['date', 'channel']
)

# Select which channels to optimize
channels_to_optimize = ['TV', 'Search']  # Only optimize these 2 channels

# Create the budget mask (removing date dimension)
# First, take the max across dates to get a template with only channel dimension
channel_template = channel_data.max('date')  # Shape: (channel)

# Create boolean mask - True only for channels we want to optimize
budgets_to_optimize = xr.DataArray(
    [ch in channels_to_optimize for ch in channels],
    coords={'channel': channels},
    dims=['channel']
)

print("Channel data shape:", channel_data.shape)
print("Budgets to optimize mask:")
print(budgets_to_optimize)

[cetagostini]

@TeemuSailynoja Are you planning to add the example? if so, we can merge after otherwise, we can just add an issue and merge this one!

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 91.59%. Comparing base (4c4f251) to head (8610b46).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1705   +/-   ##
=======================================
  Coverage   91.59%   91.59%           
=======================================
  Files          60       60           
  Lines        6778     6778           
=======================================
  Hits         6208     6208           
  Misses        570      570           

☔ 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.