Add optimize_grid() function to the OptimizationGrid dataclass.

PiperOrigin-RevId: 734413174

Author: viktoriias

CHANGELOG.md

@@ -23,6 +23,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+* Add `optimize_grid()` function to the `OptimizationGrid` dataclass.
+
 ## [1.0.5] - 2025-03-06
 
 * Add technical support for python 3.10.

meridian/analysis/optimizer.py

@@ -102,6 +102,96 @@ def spend_step_size(self) -> float:
     """The spend step size."""
     return self.grid_dataset.attrs[c.SPEND_STEP_SIZE]
 
+  def optimize_grid(
+      self,
+      budget: float | None = None,
+      fixed_budget: bool = True,
+      target_mroi: float | None = None,
+      target_roi: float | None = None,
+  ) -> np.ndarray:
+    """Hill-climbing search algorithm for budget optimization.
+
+    Args:
+      budget: Number indicating the total budget for the fixed budget scenario.
+        Defaults to the budget used in the grid.
+      fixed_budget: Bool indicating whether it's a fixed budget optimization or
+        flexible budget optimization.
+      target_mroi: Optional float indicating the target marginal return on
+        investment (mroi) constraint. This can be translated into "How much can
+        I spend when I have flexible budget until the mroi of each channel hits
+        the target mroi". It's still possible that the mroi of some channels
+        will not be equal to the target mroi due to the feasible range of media
+        spend. However, the mroi will effectively shrink toward the target mroi.
+      target_roi: Optional float indicating the target return on investment
+        (roi) constraint. This can be translated into "How much can I spend when
+        I have a flexible budget until the roi of total media spend hits the
+        target roi".
+
+    Returns:
+      optimal_spend: np.ndarry of dimension `(n_total_channels,)` containing the
+        media spend that maximizes incremental outcome based on spend
+        constraints for all media and RF channels.
+    """
+    _validate_budget(
+        fixed_budget=fixed_budget,
+        budget=budget,
+        target_mroi=target_mroi,
+        target_roi=target_roi,
+    )
+    if fixed_budget and budget is None:
+      rounded_spend = np.round(self.spend, self.round_factor).astype(int)
+      budget = np.sum(rounded_spend)
+
+    spend = self.spend_grid[0, :].copy()
+    incremental_outcome = self.incremental_outcome_grid[0, :].copy()
+    spend_grid = self.spend_grid[1:, :]
+    incremental_outcome_grid = self.incremental_outcome_grid[1:, :]
+    iterative_roi_grid = np.round(
+        tf.math.divide_no_nan(
+            incremental_outcome_grid - incremental_outcome, spend_grid - spend
+        ),
+        decimals=8,
+    )
+    while True:
+      spend_optimal = spend.astype(int)
+      # If none of the exit criteria are met roi_grid will eventually be filled
+      # with all nans.
+      if np.isnan(iterative_roi_grid).all():
+        break
+      point = np.unravel_index(
+          np.nanargmax(iterative_roi_grid), iterative_roi_grid.shape
+      )
+      row_idx = point[0]
+      media_idx = point[1]
+      spend[media_idx] = spend_grid[row_idx, media_idx]
+      incremental_outcome[media_idx] = incremental_outcome_grid[
+          row_idx, media_idx
+      ]
+      roi_grid_point = iterative_roi_grid[row_idx, media_idx]
+      if _exceeds_optimization_constraints(
+          fixed_budget=fixed_budget,
+          budget=budget,
+          spend=spend,
+          incremental_outcome=incremental_outcome,
+          roi_grid_point=roi_grid_point,
+          target_mroi=target_mroi,
+          target_roi=target_roi,
+      ):
+        break
+
+      iterative_roi_grid[0 : row_idx + 1, media_idx] = np.nan
+      iterative_roi_grid[row_idx + 1 :, media_idx] = np.round(
+          tf.math.divide_no_nan(
+              incremental_outcome_grid[row_idx + 1 :, media_idx]
+              - incremental_outcome_grid[row_idx, media_idx],
+              spend_grid[row_idx + 1 :, media_idx]
+              - spend_grid[row_idx, media_idx],
+          ),
+          decimals=8,
+      )
+
+    return spend_optimal
+
 
 @dataclasses.dataclass(frozen=True)
 class OptimizationResults:
@@ -1051,6 +1141,9 @@ def optimize(
         include_rf=self._meridian.n_rf_channels > 0,
     ).data
 
+    use_historical_budget = budget is None or round(budget) == round(
+        np.sum(hist_spend)
+    )
     budget = budget or np.sum(hist_spend)
     pct_of_spend = self._validate_pct_of_spend(hist_spend, pct_of_spend)
     spend = budget * pct_of_spend
@@ -1089,18 +1182,16 @@ def optimize(
         optimal_frequency=optimal_frequency,
         batch_size=batch_size,
     )
-    # TODO: b/375644691) - Move grid search to a OptimizationGrid class.
-    optimal_spend = self._grid_search(
-        spend_grid=optimization_grid.spend_grid,
-        incremental_outcome_grid=optimization_grid.incremental_outcome_grid,
-        budget=np.sum(rounded_spend),
+    if use_historical_budget:
+      optimization_budget = None
+    else:
+      optimization_budget = np.sum(rounded_spend)
+    optimal_spend = optimization_grid.optimize_grid(
+        budget=optimization_budget,
         fixed_budget=fixed_budget,
         target_mroi=target_mroi,
         target_roi=target_roi,
     )
-    use_historical_budget = budget is None or round(budget) == round(
-        np.sum(hist_spend)
-    )
     nonoptimized_data = self._create_budget_dataset(
         use_posterior=use_posterior,
         use_kpi=use_kpi,
@@ -1168,7 +1259,7 @@ def create_optimization_grid(
       round_factor: int,
       use_posterior: bool = True,
       use_kpi: bool = False,
-      use_optimal_frequency: bool = True,
+      use_optimal_frequency: bool = False,
       optimal_frequency: xr.DataArray | None = None,
       batch_size: int = c.DEFAULT_BATCH_SIZE,
   ) -> OptimizationGrid:
@@ -1810,96 +1901,6 @@ def _create_grids(
       )
     return (spend_grid, incremental_outcome_grid)
 
-  def _grid_search(
-      self,
-      spend_grid: np.ndarray,
-      incremental_outcome_grid: np.ndarray,
-      budget: float,
-      fixed_budget: bool,
-      target_mroi: float | None = None,
-      target_roi: float | None = None,
-  ) -> np.ndarray:
-    """Hill-climbing search algorithm for budget optimization.
-
-    Args:
-      spend_grid: Discrete grid with dimensions (`grid_length` x
-        `n_total_channels`) containing spend by channel for all media and RF
-        channels, used in the hill-climbing search algorithm.
-      incremental_outcome_grid: Discrete grid with dimensions (`grid_length` x
-        `n_total_channels`) containing incremental outcome by channel for all
-        media and RF channels, used in the hill-climbing search algorithm.
-      budget: Integer indicating the total budget.
-      fixed_budget: Bool indicating whether it's a fixed budget optimization or
-        flexible budget optimization.
-      target_mroi: Optional float indicating the target marginal return on
-        investment (mroi) constraint. This can be translated into "How much can
-        I spend when I have flexible budget until the mroi of each channel hits
-        the target mroi". It's still possible that the mroi of some channels
-        will not be equal to the target mroi due to the feasible range of media
-        spend. However, the mroi will effectively shrink toward the target mroi.
-      target_roi: Optional float indicating the target return on investment
-        (roi) constraint. This can be translated into "How much can I spend when
-        I have a flexible budget until the roi of total media spend hits the
-        target roi".
-
-    Returns:
-      optimal_spend: np.ndarry of dimension (`n_total_channels`) containing the
-        media spend that maximizes incremental outcome based on spend
-        constraints for all media and RF channels.
-      optimal_inc_outcome: np.ndarry of dimension (`n_total_channels`)
-        containing the post optimization incremental outcome per channel for all
-        media and RF channels.
-    """
-    spend = spend_grid[0, :].copy()
-    incremental_outcome = incremental_outcome_grid[0, :].copy()
-    spend_grid = spend_grid[1:, :]
-    incremental_outcome_grid = incremental_outcome_grid[1:, :]
-    iterative_roi_grid = np.round(
-        tf.math.divide_no_nan(
-            incremental_outcome_grid - incremental_outcome, spend_grid - spend
-        ),
-        decimals=8,
-    )
-    while True:
-      spend_optimal = spend.astype(int)
-      # If none of the exit criteria are met roi_grid will eventually be filled
-      # with all nans.
-      if np.isnan(iterative_roi_grid).all():
-        break
-      point = np.unravel_index(
-          np.nanargmax(iterative_roi_grid), iterative_roi_grid.shape
-      )
-      row_idx = point[0]
-      media_idx = point[1]
-      spend[media_idx] = spend_grid[row_idx, media_idx]
-      incremental_outcome[media_idx] = incremental_outcome_grid[
-          row_idx, media_idx
-      ]
-      roi_grid_point = iterative_roi_grid[row_idx, media_idx]
-      if _exceeds_optimization_constraints(
-          fixed_budget,
-          budget,
-          spend,
-          incremental_outcome,
-          roi_grid_point,
-          target_mroi,
-          target_roi,
-      ):
-        break
-
-      iterative_roi_grid[0 : row_idx + 1, media_idx] = np.nan
-      iterative_roi_grid[row_idx + 1 :, media_idx] = np.round(
-          tf.math.divide_no_nan(
-              incremental_outcome_grid[row_idx + 1 :, media_idx]
-              - incremental_outcome_grid[row_idx, media_idx],
-              spend_grid[row_idx + 1 :, media_idx]
-              - spend_grid[row_idx, media_idx],
-          ),
-          decimals=8,
-      )
-    return spend_optimal
-
-
 def _validate_budget(
     fixed_budget: bool,
     budget: float | None,