google-meridian 1.0.7__py3-none-any.whl → 1.0.8__py3-none-any.whl

meridian/analysis/optimizer.py

@@ -37,6 +37,7 @@ import xarray as xr
 
 __all__ = [
     'BudgetOptimizer',
+    'OptimizationGrid',
     'OptimizationResults',
 ]
 
@@ -92,10 +93,14 @@ class OptimizationGrid:
   Attributes:
     historical_spend: ndarray of shape `(n_paid_channels,)` containing
       aggregated historical spend allocation for spend for all media and RF
-      channels. The order matches `InputData.get_all_paid_channels`.
+      channels.
     use_kpi: Whether using generic KPI or revenue.
     use_posterior: Whether posterior distributions were used, or prior.
     use_optimal_frequency: Whether optimal frequency was used.
+    gtol: Float indicating the acceptable relative error for the budget used in
+      the grid setup. The budget is rounded by `10*n`, where `n` is the smallest
+      integer such that `(budget - rounded_budget)` is less than or equal to
+      `(budget * gtol)`.
     round_factor: The round factor used for the optimization grid.
     optimal_frequency: Optional ndarray of shape `(n_paid_channels,)`,
       containing the optimal frequency per channel. Value is `None` if the model
@@ -111,6 +116,7 @@ class OptimizationGrid:
   use_kpi: bool
   use_posterior: bool
   use_optimal_frequency: bool
+  gtol: float
   round_factor: int
   optimal_frequency: np.ndarray | None
   selected_times: list[str] | None
@@ -142,35 +148,149 @@ class OptimizationGrid:
     """The spend step size."""
     return self.grid_dataset.attrs[c.SPEND_STEP_SIZE]
 
-  # TODO: b/402950014 - Add per-channel constraints parameter.
+  @property
+  def channels(self) -> list[str]:
+    """The spend channels in the grid."""
+    return self.grid_dataset.channel.data.tolist()
+
   def optimize(
       self,
       scenario: FixedBudgetScenario | FlexibleBudgetScenario,
+      pct_of_spend: Sequence[float] | None = None,
+      spend_constraint_lower: _SpendConstraint | None = None,
+      spend_constraint_upper: _SpendConstraint | None = None,
+  ) -> xr.Dataset:
+    """Finds the optimal budget allocation that maximizes outcome.
+
+    Args:
+      scenario: The optimization scenario with corresponding parameters.
+      pct_of_spend: Numeric list of size `channels` containing the percentage
+        allocation for spend for all channels. The values must be between 0-1,
+        summing to 1. By default, the historical allocation is used. Budget and
+        allocation are used in conjunction to determine the non-optimized
+        media-level spend, which is used to calculate the non-optimized
+        performance metrics (for example, ROI) and construct the feasible range
+        of media-level spend with the spend constraints.
+      spend_constraint_lower: Numeric list of size `channels` or float (same
+        constraint for all channels) indicating the lower bound of media-level
+        spend. If given as a channel-indexed array, the order must match
+        `channels`. The lower bound of media-level spend is `(1 -
+        spend_constraint_lower) * budget * allocation)`. The value must be
+        between 0-1. Defaults to `0.3` for fixed budget and `1` for flexible.
+      spend_constraint_upper: Numeric list of size `channels` or float (same
+        constraint for all channels) indicating the upper bound of media-level
+        spend. If given as a channel-indexed array, the order must match
+        `channels`. The upper bound of media-level spend is `(1 +
+        spend_constraint_upper) * budget * allocation)`. Defaults to `0.3` for
+        fixed budget and `1` for flexible.
+
+    Returns:
+      An xarray Dataset with `channel` as the coordinate and the following data
+      variables:
+        * `optimized`: media spend that maximizes incremental outcome based
+        on spend constraints for all media and RF channels.
+        * `non_optimized`: Channel-level spend.
+
+    Raises:
+      A warning if the budget's rounding should be different from the grid's
+      round factor.'.
+      ValueError: If spend allocation is not within the grid coverage.
+    """
+    total_budget = (
+        scenario.total_budget
+        if isinstance(scenario, FixedBudgetScenario)
+        else None
+    )
+    budget = total_budget or np.sum(self.historical_spend)
+    valid_pct_of_spend = _validate_pct_of_spend(
+        n_channels=len(self.channels),
+        hist_spend=self.historical_spend,
+        pct_of_spend=pct_of_spend,
+    )
+    spend = budget * valid_pct_of_spend
+    spend_constraint_default = (
+        c.SPEND_CONSTRAINT_DEFAULT_FIXED_BUDGET
+        if isinstance(scenario, FixedBudgetScenario)
+        else c.SPEND_CONSTRAINT_DEFAULT_FLEXIBLE_BUDGET
+    )
+    if spend_constraint_lower is None:
+      spend_constraint_lower = spend_constraint_default
+    if spend_constraint_upper is None:
+      spend_constraint_upper = spend_constraint_default
+    (optimization_lower_bound, optimization_upper_bound) = (
+        _get_optimization_bounds(
+            n_channels=len(self.channels),
+            spend=spend,
+            round_factor=self.round_factor,
+            spend_constraint_lower=spend_constraint_lower,
+            spend_constraint_upper=spend_constraint_upper,
+        )
+    )
+    self._check_optimization_bounds(
+        lower_bound=optimization_lower_bound,
+        upper_bound=optimization_upper_bound,
+    )
+    round_factor = _get_round_factor(budget, self.gtol)
+    if round_factor != self.round_factor:
+      warnings.warn(
+          'Optimization accuracy may suffer owing to budget level differences.'
+          ' Consider creating a new grid with smaller `gtol` if you intend to'
+          " shrink budgets significantly. It's only a problem when you use a"
+          ' smaller budget, for which the intended step size is meant to be'
+          ' smaller for one or more channels.'
+      )
+    (spend_grid, incremental_outcome_grid) = self._trim_grid(
+        spend_bound_lower=optimization_lower_bound,
+        spend_bound_upper=optimization_upper_bound,
+    )
+    if isinstance(scenario, FixedBudgetScenario):
+      rounded_spend = np.round(spend, self.round_factor)
+      scenario = dataclasses.replace(
+          scenario, total_budget=np.sum(rounded_spend)
+      )
+    optimal_spend = self._grid_search(
+        spend_grid=spend_grid,
+        incremental_outcome_grid=incremental_outcome_grid,
+        scenario=scenario,
+    )
+
+    return xr.Dataset(
+        coords={c.CHANNEL: self.channels},
+        data_vars={
+            c.OPTIMIZED: ([c.CHANNEL], optimal_spend.data),
+            c.NON_OPTIMIZED: ([c.CHANNEL], spend),
+        },
+    )
+
+  def _grid_search(
+      self,
+      spend_grid: np.ndarray,
+      incremental_outcome_grid: np.ndarray,
+      scenario: FixedBudgetScenario | FlexibleBudgetScenario,
   ) -> 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.
       scenario: The optimization scenario with corresponding parameters.
 
     Returns:
-      optimal_spend: `np.ndarray` with shape `(n_paid_channels,)` containing the
-        media spend that maximizes incremental outcome based on spend
+      optimal_spend: `np.ndarray` 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.ndarray` of dimension (`n_total_channels`)
+        containing the post optimization incremental outcome per channel for all
+        media and RF channels.
     """
-    if (
-        isinstance(scenario, FixedBudgetScenario)
-        and scenario.total_budget is None
-    ):
-      rounded_spend = np.round(self.historical_spend, self.round_factor).astype(
-          int
-      )
-      budget = np.sum(rounded_spend)
-      scenario = dataclasses.replace(scenario, total_budget=budget)
-
-    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:, :]
+    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
@@ -211,9 +331,97 @@ class OptimizationGrid:
           ),
           decimals=8,
       )
-
     return spend_optimal
 
+  def _trim_grid(
+      self,
+      spend_bound_lower: np.ndarray,
+      spend_bound_upper: np.ndarray,
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Trim the grids based on a more restricted spend bound.
+
+    It is assumed that spend bounds are validated: their values are within the
+    grid coverage and they are rounded using this grid's round factor.
+
+    Args:
+      spend_bound_lower: The lower bound of spend for each channel.
+      spend_bound_upper: The upper bound of spend for each channel.
+
+    Returns:
+      updated_spend: The updated spend grid with valid spend values moved up to
+        the first row and invalid spend values filled with NaN.
+      updated_incremental_outcome: The updated incremental outcome grid with the
+        corresponding incremental outcome values moved up to the first row and
+        invalid incremental outcome values filled with NaN.
+    """
+    spend_grid = self.spend_grid
+    updated_spend = self.spend_grid.copy()
+    updated_incremental_outcome = self.incremental_outcome_grid.copy()
+
+    for ch in range(len(self.channels)):
+      valid_indices = np.where(
+          (spend_grid[:, ch] >= spend_bound_lower[ch])
+          & (spend_grid[:, ch] <= spend_bound_upper[ch])
+      )[0]
+      first_valid_index = valid_indices[0]
+      last_valid_index = valid_indices[-1]
+
+      # Move the smallest spend to the first row.
+      updated_spend[:, ch] = np.roll(
+          updated_spend[:, ch], shift=-first_valid_index
+      )
+      # Move the corresponding incremental outcome to the first row.
+      updated_incremental_outcome[:, ch] = np.roll(
+          updated_incremental_outcome[:, ch], shift=-first_valid_index
+      )
+
+      # Fill the invalid indices with NaN.
+      nan_indices = last_valid_index - first_valid_index + 1
+      updated_spend[nan_indices:, ch] = np.nan
+      updated_incremental_outcome[nan_indices:, ch] = np.nan
+
+    return (updated_spend, updated_incremental_outcome)
+
+  def _check_optimization_bounds(
+      self,
+      lower_bound: np.ndarray,
+      upper_bound: np.ndarray,
+  ) -> None:
+    """Checks if the spend grid fits within the optimization bounds.
+
+    Args:
+      lower_bound: `np.ndarray` of shape `(n_channels,)` containing the lower
+        bound for each channel.
+      upper_bound: `np.ndarray` of shape `(n_channels,)` containing the upper
+        bound for each channel.
+
+    Raises:
+      ValueError: If the spend grid does not fit within the optimization bounds.
+    """
+    min_spend = np.min(self.spend_grid, axis=0)
+    max_spend = np.max(self.spend_grid, axis=0)
+    errors = []
+    for i, channel_min_spend in enumerate(min_spend.data):
+      if lower_bound[i] < channel_min_spend:
+        errors.append(
+            f'Lower bound {lower_bound[i]} for channel'
+            f' {self.channels[i]} is below the mimimum spend of the grid'
+            f' {channel_min_spend}.'
+        )
+    for i, channel_max_spend in enumerate(max_spend.data):
+      if upper_bound[i] > channel_max_spend:
+        errors.append(
+            f'Upper bound {upper_bound[i]} for channel'
+            f' {self.channels[i]} is above the maximum spend of the grid'
+            f' {channel_max_spend}.'
+        )
+
+    if errors:
+      raise ValueError(
+          'Spend allocation is not within the grid coverage:\n'
+          + '\n'.join(errors)
+      )
+
 
 @dataclasses.dataclass(frozen=True)
 class OptimizationResults:
@@ -490,7 +698,7 @@ class OptimizationResults:
             title=formatter.custom_title_params(
                 summary_text.SPEND_ALLOCATION_CHART_TITLE
             ),
-            width=c.VEGALITE_FACET_DEFAULT_WIDTH
+            width=c.VEGALITE_FACET_DEFAULT_WIDTH,
         )
     )
 
@@ -698,6 +906,7 @@ class OptimizationResults:
         use_posterior=self.optimization_grid.use_posterior,
         selected_times=selected_times,
         by_reach=True,
+        use_kpi=not self.nonoptimized_data.attrs[c.IS_REVENUE_KPI],
         use_optimal_frequency=self.optimization_grid.use_optimal_frequency,
     )
 
@@ -1149,66 +1358,29 @@ class BudgetOptimizer:
         target_roi=target_roi,
         target_mroi=target_mroi,
     )
-
-    if selected_times is not None:
-      start_date, end_date = selected_times
-      selected_time_dims = self._meridian.expand_selected_time_dims(
-          start_date=start_date,
-          end_date=end_date,
-      )
-    else:
-      selected_time_dims = None
-    hist_spend = self._analyzer.get_historical_spend(
-        selected_time_dims,
-        include_media=self._meridian.n_media_channels > 0,
-        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
-    round_factor = _get_round_factor(budget, gtol)
-    rounded_spend = np.round(spend, round_factor).astype(int)
-    if self._meridian.n_rf_channels > 0 and use_optimal_frequency:
-      optimal_frequency = tf.convert_to_tensor(
-          self._analyzer.optimal_freq(
-              use_posterior=use_posterior,
-              selected_times=selected_time_dims,
-              use_kpi=use_kpi,
-          ).optimal_frequency,
-          dtype=tf.float32,
-      )
-    else:
-      optimal_frequency = None
-
-    (optimization_lower_bound, optimization_upper_bound, spend_bounds) = (
-        self._get_optimization_bounds(
-            spend=rounded_spend,
-            spend_constraint_lower=spend_constraint_lower,
-            spend_constraint_upper=spend_constraint_upper,
-            round_factor=round_factor,
-            fixed_budget=fixed_budget,
-        )
+    spend_constraint_default = (
+        c.SPEND_CONSTRAINT_DEFAULT_FIXED_BUDGET
+        if fixed_budget
+        else c.SPEND_CONSTRAINT_DEFAULT_FLEXIBLE_BUDGET
     )
+    if spend_constraint_lower is None:
+      spend_constraint_lower = spend_constraint_default
+    if spend_constraint_upper is None:
+      spend_constraint_upper = spend_constraint_default
     optimization_grid = self.create_optimization_grid(
-        historical_spend=hist_spend,
-        spend_bound_lower=optimization_lower_bound,
-        spend_bound_upper=optimization_upper_bound,
-        selected_times=selected_time_dims,
-        round_factor=round_factor,
+        selected_times=selected_times,
+        budget=budget,
+        pct_of_spend=pct_of_spend,
+        spend_constraint_lower=spend_constraint_lower,
+        spend_constraint_upper=spend_constraint_upper,
+        gtol=gtol,
         use_posterior=use_posterior,
         use_kpi=use_kpi,
         use_optimal_frequency=use_optimal_frequency,
-        optimal_frequency=optimal_frequency,
         batch_size=batch_size,
     )
-
     if fixed_budget:
-      total_budget = None if use_historical_budget else np.sum(rounded_spend)
-      scenario = FixedBudgetScenario(total_budget=total_budget)
+      scenario = FixedBudgetScenario(total_budget=budget)
     elif target_roi:
       scenario = FlexibleBudgetScenario(
           target_metric=c.ROI, target_value=target_roi
@@ -1217,16 +1389,25 @@ class BudgetOptimizer:
       scenario = FlexibleBudgetScenario(
           target_metric=c.MROI, target_value=target_mroi
       )
-
-    optimal_spend = optimization_grid.optimize(
+    spend = optimization_grid.optimize(
         scenario=scenario,
+        pct_of_spend=pct_of_spend,
+        spend_constraint_lower=spend_constraint_lower,
+        spend_constraint_upper=spend_constraint_upper,
+    )
+
+    use_historical_budget = budget is None or np.isclose(
+        budget, np.sum(optimization_grid.historical_spend)
     )
+    rounded_spend = np.round(
+        spend.non_optimized, optimization_grid.round_factor
+    ).astype(int)
     nonoptimized_data = self._create_budget_dataset(
         use_posterior=use_posterior,
         use_kpi=use_kpi,
-        hist_spend=hist_spend,
+        hist_spend=optimization_grid.historical_spend,
         spend=rounded_spend,
-        selected_times=selected_time_dims,
+        selected_times=optimization_grid.selected_times,
         confidence_level=confidence_level,
         batch_size=batch_size,
         use_historical_budget=use_historical_budget,
@@ -1234,10 +1415,10 @@ class BudgetOptimizer:
     nonoptimized_data_with_optimal_freq = self._create_budget_dataset(
         use_posterior=use_posterior,
         use_kpi=use_kpi,
-        hist_spend=hist_spend,
+        hist_spend=optimization_grid.historical_spend,
         spend=rounded_spend,
-        selected_times=selected_time_dims,
-        optimal_frequency=optimal_frequency,
+        selected_times=optimization_grid.selected_times,
+        optimal_frequency=optimization_grid.optimal_frequency,
         confidence_level=confidence_level,
         batch_size=batch_size,
         use_historical_budget=use_historical_budget,
@@ -1252,10 +1433,10 @@ class BudgetOptimizer:
     optimized_data = self._create_budget_dataset(
         use_posterior=use_posterior,
         use_kpi=use_kpi,
-        hist_spend=hist_spend,
-        spend=optimal_spend,
-        selected_times=selected_time_dims,
-        optimal_frequency=optimal_frequency,
+        hist_spend=optimization_grid.historical_spend,
+        spend=spend.optimized,
+        selected_times=optimization_grid.selected_times,
+        optimal_frequency=optimization_grid.optimal_frequency,
         attrs=constraints,
         confidence_level=confidence_level,
         batch_size=batch_size,
@@ -1263,17 +1444,23 @@ class BudgetOptimizer:
     )
 
     if not fixed_budget:
-      self._raise_warning_if_target_constraints_not_met(
+      _raise_warning_if_target_constraints_not_met(
           target_roi=target_roi,
           target_mroi=target_mroi,
           optimized_data=optimized_data,
       )
 
     spend_ratio = np.divide(
-        spend,
-        hist_spend,
-        out=np.zeros_like(hist_spend, dtype=float),
-        where=hist_spend != 0,
+        spend.non_optimized,
+        optimization_grid.historical_spend,
+        out=np.zeros_like(optimization_grid.historical_spend, dtype=float),
+        where=optimization_grid.historical_spend != 0,
+    )
+    n_paid_channels = len(self._meridian.input_data.get_all_paid_channels())
+    spend_bounds = _get_spend_bounds(
+        n_channels=n_paid_channels,
+        spend_constraint_lower=spend_constraint_lower,
+        spend_constraint_upper=spend_constraint_upper,
     )
 
     return OptimizationResults(
@@ -1287,71 +1474,68 @@ class BudgetOptimizer:
         _optimization_grid=optimization_grid,
     )
 
-  def _raise_warning_if_target_constraints_not_met(
-      self,
-      target_roi: float | None,
-      target_mroi: float | None,
-      optimized_data: xr.Dataset,
-  ) -> None:
-    """Raises a warning if the target constraints are not met."""
-    if target_roi:
-      # Total ROI is a scalar value.
-      optimized_roi = optimized_data.attrs[c.TOTAL_ROI]
-      if optimized_roi < target_roi:
-        warnings.warn(
-            f'Target ROI constraint was not met. The target ROI is {target_roi}'
-            f', but the actual ROI is {optimized_roi}.'
-        )
-    elif target_mroi:
-      # Compare each channel's marginal ROI to the target.
-      # optimized_data[c.MROI] is an array of shape (n_channels, 4), where the
-      # last dimension is [mean, median, ci_lo, ci_hi].
-      optimized_mroi = optimized_data[c.MROI][:, 0]
-      if np.any(optimized_mroi < target_mroi):
-        warnings.warn(
-            'Target marginal ROI constraint was not met. The target marginal'
-            f' ROI is {target_mroi}, but the actual channel marginal ROIs are'
-            f' {optimized_mroi}.'
-        )
-
   def create_optimization_grid(
       self,
-      historical_spend: np.ndarray,
-      spend_bound_lower: np.ndarray,
-      spend_bound_upper: np.ndarray,
-      selected_times: Sequence[str] | None,
-      round_factor: int,
       use_posterior: bool = True,
+      selected_times: tuple[str | None, str | None] | None = None,
+      budget: float | None = None,
+      pct_of_spend: Sequence[float] | None = None,
+      spend_constraint_lower: _SpendConstraint = c.SPEND_CONSTRAINT_DEFAULT,
+      spend_constraint_upper: _SpendConstraint = c.SPEND_CONSTRAINT_DEFAULT,
+      gtol: float = 0.0001,
+      use_optimal_frequency: bool = True,
       use_kpi: bool = False,
-      use_optimal_frequency: bool = False,
-      optimal_frequency: xr.DataArray | None = None,
       batch_size: int = c.DEFAULT_BATCH_SIZE,
   ) -> OptimizationGrid:
     """Creates a OptimizationGrid for optimization.
 
     Args:
-      historical_spend: ndarray of shape `(n_paid_channels,)` with arrgegated
-        historical spend per paid channel.
-      spend_bound_lower: ndarray of dimension `(n_total_channels,)` containing
-        the lower constraint spend for each channel.
-      spend_bound_upper: ndarray of dimension `(n_total_channels,)` containing
-        the upper constraint spend for each channel.
-      selected_times: Sequence of strings representing the time dimensions in
-        `meridian.input_data.time` to use for optimization.
-      round_factor: The round factor used for the optimization grid.
       use_posterior: Boolean. If `True`, then the incremental outcome is derived
         from the posterior distribution of the model. Otherwise, the prior
         distribution is used.
+      selected_times: Tuple containing the start and end time dimension
+        coordinates for the duration to run the optimization on. Selected time
+        values should align with the Meridian time dimension coordinates in the
+        underlying model. By default, all times periods are used. Either start
+        or end time component can be `None` to represent the first or the last
+        time coordinate, respectively.
+      budget: Number indicating the total budget for the fixed budget scenario.
+        Defaults to the historical budget.
+      pct_of_spend: Numeric list of size `n_paid_channels` containing the
+        percentage allocation for spend for all media and RF channels. The order
+        must match `(InputData.media + InputData.reach)` with values between
+        0-1, summing to 1. By default, the historical allocation is used. Budget
+        and allocation are used in conjunction to determine the non-optimized
+        media-level spend, which is used to calculate the non-optimized
+        performance metrics (for example, ROI) and construct the feasible range
+        of media-level spend with the spend constraints. Consider using
+        `InputData.get_paid_channels_argument_builder()` to construct this
+        argument.
+      spend_constraint_lower: Numeric list of size `n_paid_channels` or float
+        (same constraint for all channels) indicating the lower bound of
+        media-level spend. If given as a channel-indexed array, the order must
+        match `(InputData.media + InputData.reach)`. The lower bound of
+        media-level spend is `(1 - spend_constraint_lower) * budget *
+        allocation)`. The value must be between 0-1. Defaults to `0.3` for fixed
+        budget and `1` for flexible. Consider using
+        `InputData.get_paid_channels_argument_builder()` to construct this
+        argument.
+      spend_constraint_upper: Numeric list of size `n_paid_channels` or float
+        (same constraint for all channels) indicating the upper bound of
+        media-level spend. If given as a channel-indexed array, the order must
+        match `(InputData.media + InputData.reach)`. The upper bound of
+        media-level spend is `(1 + spend_constraint_upper) * budget *
+        allocation)`. Defaults to `0.3` for fixed budget and `1` for flexible.
+        Consider using `InputData.get_paid_channels_argument_builder()` to
+        construct this argument.
+      gtol: Float indicating the acceptable relative error for the budget used
+        in the grid setup. The budget will be rounded by `10*n`, where `n` is
+        the smallest integer such that `(budget - rounded_budget)` is less than
+        or equal to `(budget * gtol)`. `gtol` must be less than 1.
+      use_optimal_frequency: Boolean. Whether optimal frequency was used.
       use_kpi: Boolean. If `True`, then the incremental outcome is derived from
         the KPI impact. Otherwise, the incremental outcome is derived from the
         revenue impact.
-      use_optimal_frequency: Boolean. Whether optimal frequency was used.
-      optimal_frequency: `xr.DataArray` with dimension `n_rf_channels`,
-        containing the optimal frequency per channel, that maximizes mean ROI
-        over the corresponding prior/posterior distribution. Value is `None` if
-        the model does not contain reach and frequency data, or if the model
-        does contain reach and frequency data, but historical frequency is used
-        for the optimization scenario.
       batch_size: Max draws per chain in each batch. The calculation is run in
         batches to avoid memory exhaustion. If a memory error occurs, try
         reducing `batch_size`. The calculation will generally be faster with
@@ -1361,14 +1545,56 @@ class BudgetOptimizer:
       An OptimizationGrid object containing the grid data for optimization.
     """
     self._validate_model_fit(use_posterior)
+    if selected_times is not None:
+      start_date, end_date = selected_times
+      selected_time_dims = self._meridian.expand_selected_time_dims(
+          start_date=start_date,
+          end_date=end_date,
+      )
+    else:
+      selected_time_dims = None
+    hist_spend = self._analyzer.get_historical_spend(
+        selected_time_dims,
+        include_media=self._meridian.n_media_channels > 0,
+        include_rf=self._meridian.n_rf_channels > 0,
+    ).data
+    n_paid_channels = len(self._meridian.input_data.get_all_paid_channels())
+    budget = budget or np.sum(hist_spend)
+    valid_pct_of_spend = _validate_pct_of_spend(
+        n_channels=n_paid_channels,
+        hist_spend=hist_spend,
+        pct_of_spend=pct_of_spend,
+    )
+    spend = budget * valid_pct_of_spend
+    round_factor = _get_round_factor(budget, gtol)
+    (optimization_lower_bound, optimization_upper_bound) = (
+        _get_optimization_bounds(
+            n_channels=n_paid_channels,
+            spend=spend,
+            round_factor=round_factor,
+            spend_constraint_lower=spend_constraint_lower,
+            spend_constraint_upper=spend_constraint_upper,
+        )
+    )
+    if self._meridian.n_rf_channels > 0 and use_optimal_frequency:
+      optimal_frequency = tf.convert_to_tensor(
+          self._analyzer.optimal_freq(
+              use_posterior=use_posterior,
+              selected_times=selected_time_dims,
+              use_kpi=use_kpi,
+          ).optimal_frequency,
+          dtype=tf.float32,
+      )
+    else:
+      optimal_frequency = None
 
     step_size = 10 ** (-round_factor)
     (spend_grid, incremental_outcome_grid) = self._create_grids(
-        spend=historical_spend,
-        spend_bound_lower=spend_bound_lower,
-        spend_bound_upper=spend_bound_upper,
+        spend=hist_spend,
+        spend_bound_lower=optimization_lower_bound,
+        spend_bound_upper=optimization_upper_bound,
         step_size=step_size,
-        selected_times=selected_times,
+        selected_times=selected_time_dims,
         use_posterior=use_posterior,
         use_kpi=use_kpi,
         optimal_frequency=optimal_frequency,
@@ -1382,13 +1608,14 @@ class BudgetOptimizer:
 
     return OptimizationGrid(
         _grid_dataset=grid_dataset,
-        historical_spend=historical_spend,
+        historical_spend=hist_spend,
         use_kpi=use_kpi,
         use_posterior=use_posterior,
         use_optimal_frequency=use_optimal_frequency,
+        gtol=gtol,
         round_factor=round_factor,
         optimal_frequency=optimal_frequency,
-        selected_times=selected_times,
+        selected_times=selected_time_dims,
     )
 
   def _create_grid_dataset(
@@ -1425,78 +1652,12 @@ class BudgetOptimizer:
     return xr.Dataset(
         data_vars=data_vars,
         coords={
-            c.GRID_SPEND_INDEX: (
-                [c.GRID_SPEND_INDEX],
-                np.arange(0, len(spend_grid)),
-            ),
-            c.CHANNEL: (
-                [c.CHANNEL],
-                self._meridian.input_data.get_all_paid_channels(),
-            ),
+            c.GRID_SPEND_INDEX: np.arange(0, len(spend_grid)),
+            c.CHANNEL: self._meridian.input_data.get_all_paid_channels(),
         },
         attrs={c.SPEND_STEP_SIZE: spend_step_size},
     )
 
-  def _validate_pct_of_spend(
-      self, hist_spend: np.ndarray, pct_of_spend: Sequence[float] | None
-  ) -> np.ndarray:
-    """Validates and returns the percent of spend."""
-    if pct_of_spend is not None:
-      if len(pct_of_spend) != len(
-          self._meridian.input_data.get_all_paid_channels()
-      ):
-        raise ValueError('Percent of spend must be specified for all channels.')
-      if not math.isclose(np.sum(pct_of_spend), 1.0, abs_tol=0.001):
-        raise ValueError('Percent of spend must sum to one.')
-      return np.array(pct_of_spend)
-    else:
-      return hist_spend / np.sum(hist_spend)
-
-  def _validate_spend_constraints(
-      self,
-      fixed_budget: bool,
-      const_lower: _SpendConstraint | None,
-      const_upper: _SpendConstraint | None,
-  ) -> tuple[np.ndarray, np.ndarray]:
-    """Validates and returns the spend constraint requirements."""
-
-    def get_const_array(const: _SpendConstraint | None) -> np.ndarray:
-      if const is None:
-        const = (
-            np.array([c.SPEND_CONSTRAINT_DEFAULT_FIXED_BUDGET])
-            if fixed_budget
-            else np.array([c.SPEND_CONSTRAINT_DEFAULT_FLEXIBLE_BUDGET])
-        )
-      elif isinstance(const, (float, int)):
-        const = np.array([const])
-      else:
-        const = np.array(const)
-      return const
-
-    const_lower = get_const_array(const_lower)
-    const_upper = get_const_array(const_upper)
-
-    if any(
-        len(const)
-        not in (1, len(self._meridian.input_data.get_all_paid_channels()))
-        for const in [const_lower, const_upper]
-    ):
-      raise ValueError(
-          'Spend constraints must be either a single constraint or be specified'
-          ' for all channels.'
-      )
-
-    for const in const_lower:
-      if not 0.0 <= const <= 1.0:
-        raise ValueError(
-            'The lower spend constraint must be between 0 and 1 inclusive.'
-        )
-    for const in const_upper:
-      if const < 0:
-        raise ValueError('The upper spend constraint must be positive.')
-
-    return (const_lower, const_upper)
-
   def _get_incremental_outcome_tensors(
       self,
       hist_spend: np.ndarray,
@@ -1717,67 +1878,12 @@ class BudgetOptimizer:
     return xr.Dataset(
         data_vars=data_vars,
         coords={
-            c.CHANNEL: (
-                [c.CHANNEL],
-                self._meridian.input_data.get_all_paid_channels(),
-            ),
-            c.METRIC: (
-                [c.METRIC],
-                [c.MEAN, c.MEDIAN, c.CI_LO, c.CI_HI],
-            ),
+            c.CHANNEL: self._meridian.input_data.get_all_paid_channels(),
+            c.METRIC: [c.MEAN, c.MEDIAN, c.CI_LO, c.CI_HI],
         },
         attrs=attributes | (attrs or {}),
     )
 
-  def _get_optimization_bounds(
-      self,
-      spend: np.ndarray,
-      spend_constraint_lower: _SpendConstraint | None,
-      spend_constraint_upper: _SpendConstraint | None,
-      round_factor: int,
-      fixed_budget: bool,
-  ) -> tuple[np.ndarray, np.ndarray, tuple[np.ndarray, np.ndarray]]:
-    """Get optimization bounds from spend and spend constraints.
-
-    Args:
-      spend: np.ndarray with size `n_total_channels` containing media-level
-        spend for all media and RF channels.
-      spend_constraint_lower: Numeric list of size `n_total_channels` or float
-        (same constraint for all media) indicating the lower bound of
-        media-level spend. The lower bound of media-level spend is `(1 -
-        spend_constraint_lower) * budget * allocation)`. The value must be
-        between 0-1.
-      spend_constraint_upper: Numeric list of size `n_total_channels` or float
-        (same constraint for all media) indicating the upper bound of
-        media-level spend. The upper bound of media-level spend is `(1 +
-        spend_constraint_upper) * budget * allocation)`.
-      round_factor: Integer number of digits to round optimization bounds.
-      fixed_budget: Boolean indicating whether it's a fixed budget optimization
-        or flexible budget optimization.
-
-    Returns:
-      lower_bound: np.ndarray of size `n_total_channels` containing the treated
-        lower bound spend for each media and RF channel.
-      upper_bound: np.ndarray of size `n_total_channels` containing the treated
-        upper bound spend for each media and RF channel.
-      spend_bounds: tuple of np.ndarray of size `n_total_channels` containing
-        the untreated lower and upper bound spend for each media and RF channel.
-    """
-    (spend_const_lower, spend_const_upper) = self._validate_spend_constraints(
-        fixed_budget, spend_constraint_lower, spend_constraint_upper
-    )
-    spend_bounds = (
-        np.maximum((1 - spend_const_lower), 0),
-        (1 + spend_const_upper),
-    )
-
-    lower_bound = np.round(
-        (spend_bounds[0] * spend),
-        round_factor,
-    ).astype(int)
-    upper_bound = np.round(spend_bounds[1] * spend, round_factor).astype(int)
-    return (lower_bound, upper_bound, spend_bounds)
-
   def _update_incremental_outcome_grid(
       self,
       i: int,
@@ -1967,6 +2073,135 @@ class BudgetOptimizer:
     return (spend_grid, incremental_outcome_grid)
 
 
+def _validate_pct_of_spend(
+    n_channels: int,
+    hist_spend: np.ndarray,
+    pct_of_spend: Sequence[float] | None,
+) -> np.ndarray:
+  """Validates and returns the percent of spend."""
+  if pct_of_spend is not None:
+    if len(pct_of_spend) != n_channels:
+      raise ValueError('Percent of spend must be specified for all channels.')
+    if not math.isclose(np.sum(pct_of_spend), 1.0, abs_tol=0.001):
+      raise ValueError('Percent of spend must sum to one.')
+    return np.array(pct_of_spend)
+  else:
+    return hist_spend / np.sum(hist_spend)
+
+
+def _validate_spend_constraints(
+    n_channels: int,
+    const_lower: _SpendConstraint,
+    const_upper: _SpendConstraint,
+) -> tuple[np.ndarray, np.ndarray]:
+  """Validates and returns the spend constraint requirements."""
+
+  def get_const_array(const: _SpendConstraint) -> np.ndarray:
+    if isinstance(const, (float, int)):
+      const = np.array([const])
+    else:
+      const = np.array(const)
+    return const
+
+  const_lower = get_const_array(const_lower)
+  const_upper = get_const_array(const_upper)
+
+  if any(
+      len(const) not in (1, n_channels) for const in [const_lower, const_upper]
+  ):
+    raise ValueError(
+        'Spend constraints must be either a single constraint or be specified'
+        ' for all channels.'
+    )
+
+  for const in const_lower:
+    if not 0.0 <= const <= 1.0:
+      raise ValueError(
+          'The lower spend constraint must be between 0 and 1 inclusive.'
+      )
+  for const in const_upper:
+    if const < 0:
+      raise ValueError('The upper spend constraint must be positive.')
+
+  return (const_lower, const_upper)
+
+
+def _get_spend_bounds(
+    n_channels: int,
+    spend_constraint_lower: _SpendConstraint,
+    spend_constraint_upper: _SpendConstraint,
+) -> tuple[np.ndarray, np.ndarray]:
+  """Get spend bounds from spend constraints.
+
+  Args:
+    n_channels: Integer number of total channels.
+    spend_constraint_lower: Numeric list of size `n_total_channels` or float
+      (same constraint for all media) indicating the lower bound of media-level
+      spend. The lower bound of media-level spend is `(1 -
+      spend_constraint_lower) * budget * allocation)`. The value must be between
+      0-1.
+    spend_constraint_upper: Numeric list of size `n_total_channels` or float
+      (same constraint for all media) indicating the upper bound of media-level
+      spend. The upper bound of media-level spend is `(1 +
+      spend_constraint_upper) * budget * allocation)`.
+
+  Returns:
+    spend_bounds: tuple of np.ndarray of size `n_total_channels` containing
+        the untreated lower and upper bound spend for each media and RF channel.
+  """
+  (spend_const_lower, spend_const_upper) = _validate_spend_constraints(
+      n_channels,
+      spend_constraint_lower,
+      spend_constraint_upper,
+  )
+  spend_bounds = (
+      np.maximum((1 - spend_const_lower), 0),
+      (1 + spend_const_upper),
+  )
+  return spend_bounds
+
+
+def _get_optimization_bounds(
+    n_channels: int,
+    spend: np.ndarray,
+    round_factor: int,
+    spend_constraint_lower: _SpendConstraint,
+    spend_constraint_upper: _SpendConstraint,
+) -> tuple[np.ndarray, np.ndarray]:
+  """Get optimization bounds from spend and spend constraints.
+
+  Args:
+    n_channels: Integer number of total channels.
+    spend: np.ndarray with size `n_total_channels` containing media-level spend
+      for all media and RF channels.
+    round_factor: Integer number of digits to round optimization bounds.
+    spend_constraint_lower: Numeric list of size `n_total_channels` or float
+      (same constraint for all media) indicating the lower bound of media-level
+      spend. The lower bound of media-level spend is `(1 -
+      spend_constraint_lower) * budget * allocation)`. The value must be between
+      0-1.
+    spend_constraint_upper: Numeric list of size `n_total_channels` or float
+      (same constraint for all media) indicating the upper bound of media-level
+      spend. The upper bound of media-level spend is `(1 +
+      spend_constraint_upper) * budget * allocation)`.
+
+  Returns:
+    lower_bound: np.ndarray of size `n_total_channels` containing the treated
+      lower bound spend for each media and RF channel.
+    upper_bound: np.ndarray of size `n_total_channels` containing the treated
+      upper bound spend for each media and RF channel.
+  """
+  spend_bounds = _get_spend_bounds(
+      n_channels=n_channels,
+      spend_constraint_lower=spend_constraint_lower,
+      spend_constraint_upper=spend_constraint_upper,
+  )
+  rounded_spend = np.round(spend, round_factor).astype(int)
+  lower = np.round((spend_bounds[0] * rounded_spend), round_factor).astype(int)
+  upper = np.round(spend_bounds[1] * rounded_spend, round_factor).astype(int)
+  return (lower, upper)
+
+
 def _validate_budget(
     fixed_budget: bool,
     budget: float | None,
@@ -2063,3 +2298,30 @@ def _exceeds_optimization_constraints(
     return cur_total_roi < target_value and roi_grid_point < cur_total_roi
   else:
     return roi_grid_point < scenario.target_value
+
+
+def _raise_warning_if_target_constraints_not_met(
+    target_roi: float | None,
+    target_mroi: float | None,
+    optimized_data: xr.Dataset,
+) -> None:
+  """Raises a warning if the target constraints are not met."""
+  if target_roi:
+    # Total ROI is a scalar value.
+    optimized_roi = optimized_data.attrs[c.TOTAL_ROI]
+    if optimized_roi < target_roi:
+      warnings.warn(
+          f'Target ROI constraint was not met. The target ROI is {target_roi}'
+          f', but the actual ROI is {optimized_roi}.'
+      )
+  elif target_mroi:
+    # Compare each channel's marginal ROI to the target.
+    # optimized_data[c.MROI] is an array of shape (n_channels, 4), where the
+    # last dimension is [mean, median, ci_lo, ci_hi].
+    optimized_mroi = optimized_data[c.MROI][:, 0]
+    if np.any(optimized_mroi < target_mroi):
+      warnings.warn(
+          'Target marginal ROI constraint was not met. The target marginal'
+          f' ROI is {target_mroi}, but the actual channel marginal ROIs are'
+          f' {optimized_mroi}.'
+      )