Custom optimizer constraints

pymc_marketing/mmm/constraints.py

@@ -0,0 +1,66 @@
+#   Copyright 2025 The PyMC Labs Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   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.
+
+"""Constraints for the BudgetOptimizer."""
+
+import pytensor.tensor as pt
+from pymc.pytensorf import rewrite_pregrad
+from pytensor import function
+
+
+def build_default_sum_constraint(key: str = "default"):
+    """Return a constraint dict that enforces sum(budgets) == total_budget."""
+
+    def _constraint_fun(budgets_sym, total_budget_sym, optimizer):
+        return pt.sum(budgets_sym) - total_budget_sym
+
+    return dict(
+        key=key,
+        constraint_type="eq",
+        constraint_fun=_constraint_fun,
+    )  # type: ignore
+
+
+def compile_constraints_for_scipy(constraints, optimizer):
+    """Compile constraints for scipy."""
+    compiled_constraints = []
+
+    budgets = optimizer._budgets
+    budgets_flat = optimizer._budgets_flat
+    total_budget = optimizer._total_budget
+    for c in constraints.values() if isinstance(constraints, dict) else constraints:
+        ctype = c["constraint_type"]
+        sym_fun_output = c["constraint_fun"](budgets, total_budget, optimizer)
+        sym_jac_output = pt.grad(rewrite_pregrad(sym_fun_output), budgets_flat)
+
+        # Compile symbolic => python callables
+        compiled_fun = function(
+            inputs=[budgets_flat],
+            outputs=sym_fun_output,
+            on_unused_input="ignore",
+        )
+        compiled_jac = function(
+            inputs=[budgets_flat],
+            outputs=sym_jac_output,
+            on_unused_input="ignore",
+        )
+
+        compiled_constraints.append(
+            {
+                "type": ctype,
+                "fun": compiled_fun,
+                "jac": compiled_jac,
+            }
+        )
+    return compiled_constraints

pymc_marketing/mmm/mmm.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import warnings
+from collections.abc import Sequence
 from typing import Annotated, Any, Literal
 
 import arviz as az
@@ -475,6 +476,13 @@
                     dims=("date", "channel"),
                 )
 
+            # We define the deterministic variable to define the optimization by default
+            pm.Deterministic(
+                name="total_contributions",
+                var=channel_contributions.sum(axis=(-2, -1)),
+                dims=(),
+            )
+
             mu_var = intercept + channel_contributions.sum(axis=-1)
 
             if (
@@ -2228,8 +2236,10 @@
         budget: float | int,
         num_periods: int,
         budget_bounds: DataArray | dict[str, tuple[float, float]] | None = None,
-        response_variable: str = "channel_contributions",
+        response_variable: str = "total_contributions",
         utility_function: UtilityFunctionType = average_response,
+        constraints: Sequence[dict[str, Any]] = (),
+        default_constraints: bool = True,
         **minimize_kwargs,
     ) -> tuple[DataArray, OptimizeResult]:
         """Optimize the given budget based on the specified utility function over a specified time period.
@@ -2258,9 +2268,14 @@
             An xarray DataArary or dictionary specifying the lower and upper bounds for the budget allocation
             for each channel. If None, no bounds are applied.
         response_variable : str, optional
-            The response variable to optimize. Default is "channel_contributions".
+            The response variable to optimize. Default is "total_contributions".
         utility_function : UtilityFunctionType, optional
             The utility function to maximize. Default is the mean of the response distribution.
+        custom_constraints : list[dict[str, Any]], optional
+            Custom constraints for the optimization. If None, no custom constraints are applied. Format:
+            [{"key":...,"constraint_fun":...,"constraint_type":...}]
+        default_constraints : bool, optional
+            Whether to add the default sum constraint to the optimizer. Default is True.
         **minimize_kwargs
             Additional arguments to pass to the `BudgetOptimizer`.
 
@@ -2283,6 +2298,8 @@
             num_periods=num_periods,
             utility_function=utility_function,
             response_variable=response_variable,
+            custom_constraints=constraints,
+            default_constraints=default_constraints,
             model=self,
         )
 
@@ -2298,7 +2315,7 @@
         time_granularity: Literal["daily", "weekly", "monthly", "quarterly", "yearly"],
         num_periods: int,
         budget_bounds: DataArray | dict[str, tuple[float, float]] | None = None,
-        custom_constraints: dict[str, float] | None = None,
+        custom_constraints: Sequence[dict[str, Any]] | None = None,
         noise_level: float = 0.01,
         utility_function: UtilityFunctionType = average_response,
         **minimize_kwargs,
@@ -2334,7 +2351,7 @@
         budget_bounds : DatArray or dict[str, list[Any]], optional
             An xarray DataArray or a dictionary specifying the lower and upper bounds for the budget allocation
             for each channel. If None, no bounds are applied.
-        custom_constraints : dict[str, float], optional
+        custom_constraints : Sequence[dict[str, Any]], optional
             Custom constraints for the optimization. If None, no custom constraints are applied.
         noise_level : float, optional
             The level of noise added to the allocation strategy (by default 1%).
@@ -2372,12 +2389,12 @@
             model=self,
             num_periods=num_periods,
             utility_function=utility_function,
+            custom_constraints=custom_constraints,
+            default_constraints=True,
         )
-
         self.optimal_allocation, _ = allocator.allocate_budget(
             total_budget=budget,
             budget_bounds=budget_bounds,
-            custom_constraints=custom_constraints,
             **minimize_kwargs,
         )
 
@@ -2466,9 +2483,9 @@
         ax.grid(False)
         ax2.grid(False)
 
-        bars = [bars1[0], bars2[0]]
-        labels = [bar.get_label() for bar in bars]
-        ax.legend(bars, labels)
+        bars = [bars1, bars2]
+        labels = ["Allocated Spend", "Channel Contributions"]
+        ax.legend(bars, labels, loc="best")
 
         return fig, ax
 

pymc_marketing/mmm/utility.py

@@ -49,17 +49,15 @@
 UtilityFunctionType = Callable[[pt.TensorVariable, pt.TensorVariable], float]
 
 
-def _sum_non_sample_dims(samples: pt.TensorVariable) -> pt.TensorVariable:
-    """Sum any dimensions besides the first one.
-
-    This is needed until https://github.com/pymc-labs/pymc-marketing/issues/1387 is resolved.
-    """
+def _check_samples_dimensionality(samples: pt.TensorVariable) -> pt.TensorVariable:
+    """Check if samples is a 1D tensor variable."""
     ndim = samples.type.ndim
     if ndim == 1:
-        # Presumably non-sample dims were already reduced
         return samples
     else:
-        return samples.sum(axis=tuple(range(1, ndim)))
+        raise ValueError(
+            f"Function expected samples to be a 1D tensor variable. Got {ndim} dimensions."
+        )
 
 
 def _compute_quantile(x: pt.TensorVariable, q: float) -> pt.TensorVariable:
@@ -91,7 +89,7 @@
     samples: pt.TensorVariable, budgets: pt.TensorVariable
 ) -> pt.TensorVariable:
     """Compute the average response of the posterior predictive distribution."""
-    return pt.mean(_sum_non_sample_dims(samples))
+    return pt.mean(_check_samples_dimensionality(samples))
 
 
 def tail_distance(confidence_level: float = 0.75) -> UtilityFunctionType:
@@ -126,7 +124,7 @@
     def _tail_distance(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         mean = pt.mean(samples)
         q1 = _compute_quantile(samples, confidence_level)
         q2 = _compute_quantile(samples, 1 - confidence_level)
@@ -155,7 +153,7 @@
     pt.TensorVariable
         A PyTensor tensor variable representing the ROAS distribution.
     """
-    samples = _sum_non_sample_dims(samples)
+    samples = _check_samples_dimensionality(samples)
     total_budget = pt.sum(budgets)
     roas_distribution = samples / total_budget
     return roas_distribution
@@ -201,7 +199,7 @@
     def _mean_tightness_score(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         mean = pt.mean(samples)
         tail_metric = tail_distance(confidence_level)
         return mean - alpha * tail_metric(samples, budgets)
@@ -252,7 +250,7 @@
     def _value_at_risk(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         return _compute_quantile(samples, 1 - confidence_level)
 
     return _value_at_risk
@@ -301,7 +299,7 @@
     def _conditional_value_at_risk(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         VaR = _compute_quantile(samples, 1 - confidence_level)
         mask = samples <= VaR
         num_tail_losses = pt.sum(mask)
@@ -349,7 +347,7 @@
     def _sharpe_ratio(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         excess_returns = samples - risk_free_rate
         mean_excess_return = pt.mean(excess_returns)
         std_excess_return = pt.std(excess_returns, ddof=1)
@@ -397,7 +395,7 @@
     def _raroc(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         capital = pt.sum(budgets)
         expected_return = pt.mean(samples)
         risk_adjusted_return = expected_return - risk_free_rate
@@ -454,7 +452,7 @@
     def _adjusted_value_at_risk_score(
         samples: pt.TensorVariable, budgets: pt.TensorVariable
     ) -> pt.TensorVariable:
-        samples = _sum_non_sample_dims(samples)
+        samples = _check_samples_dimensionality(samples)
         var = _compute_quantile(samples, 1 - confidence_level)
         mean = pt.mean(samples)
         return (1 - risk_aversion) * mean + risk_aversion * var
@@ -563,7 +561,7 @@
     - Choueifaty, Y., & Coignard, Y. (2008). Toward Maximum Diversification. *Journal of Portfolio Management*.
     - Meucci, A. (2009). Managing Diversification. *Risk*, 22(5), 74-79.
     """
-    samples = _sum_non_sample_dims(samples)
+    samples = _check_samples_dimensionality(samples)
     weights = budgets / pt.sum(budgets)
     individual_volatilities = pt.std(samples, axis=0, ddof=1)
     cov_matrix = _covariance_matrix(samples)