How do I set custom constrains for the budget allocator?

[bart-vanvlerken]

Hi everyone, I'm using pymc-marketing version 0.15.1 and I'd like to specify different budgets based on the funnel phase like in this simple example:

• SEE: TV, Radio, OOH
• THINK: Display, Meta Prospecting
• DO: Paid Search, Meta Retargeting

I tried running the following code to add funnel level constraints, and although the output respects the constraints the resulting relative spend shares are almost identical within funnels. Can anyone spot what I'm doing wrong?

funnel_stages = {
    'awareness': ['display_spend', 'dooh_spend', 'google_ads_video_spend', 'ltv_ad_alliance_spend',
                  'ltv_ster_spend', 'ltv_talpa_spend', 'meta_awareness_spend', 'ooh_spend',
                  'pinterest_awareness_spend', 'programmatic_instream_video_spend',
                  'tiktok_awareness_spend', 'vod_spend'],
    'consideration': ['affiliate_content_spend', 'affiliate_influencer_spend',
                      'google_ads_demand_gen_spend', 'meta_consideration_spend',
                      'pinterest_consideration_spend'],
    'performance': ['affiliate_cashback_spend', 'affiliate_coupons_spend', 'affiliate_css_spend',
                    'affiliate_other_spend', 'google_ads_performance_max_spend',
                    'google_ads_search_branded_spend', 'google_ads_search_nonbranded_spend',
                    'google_ads_shopping_spend', 'meta_performance_spend',
                    'microsoft_ads_performance_max_spend', 'microsoft_ads_search_branded_spend',
                    'microsoft_ads_search_nonbranded_spend', 'microsoft_ads_shopping_spend',
                    'pinterest_performance_spend']
}

funnel_budgets = {
    'awareness': 2_000_000 / num_periods,
    'consideration': 650_000 / num_periods,
    'performance': 3_493_039 / num_periods
}

channel_pos = {ch: i for i, ch in enumerate(full_model.idata.posterior['channel'].values)}

custom_constraints = []

for funnel_stage, channels in funnel_stages.items():
    funnel_index = [channel_pos[ch] for ch in channels]
    funnel_budget = funnel_budgets[funnel_stage]

    def funnel_budget_fun(budgets_sym, total_budget_sym, optimizer, index = funnel_index, budget = funnel_budget):
        return pt.sum(pt.stack([budgets_sym[i] for i in index])) - budget

    custom_constraints.append(
        Constraint(
            key = f'{funnel_stage}_budget',
            constraint_type = 'eq',
            constraint_fun = funnel_budget_fun
        )
    )

budget_optimizer = BudgetOptimizer(
    num_periods = num_periods, 
    model = full_model,
    custom_constraints = custom_constraints
)