google-meridian 1.1.5__py3-none-any.whl → 1.2.0__py3-none-any.whl

meridian/analysis/optimizer.py

@@ -24,6 +24,7 @@ import warnings
 
 import altair as alt
 import jinja2
+from meridian import backend
 from meridian import constants as c
 from meridian.analysis import analyzer
 from meridian.analysis import formatter
@@ -32,7 +33,6 @@ from meridian.data import time_coordinates as tc
 from meridian.model import model
 import numpy as np
 import pandas as pd
-import tensorflow as tf
 import xarray as xr
 
 
@@ -40,6 +40,10 @@ __all__ = [
     'BudgetOptimizer',
     'OptimizationGrid',
     'OptimizationResults',
+    'FixedBudgetScenario',
+    'FlexibleBudgetScenario',
+    'get_optimization_bounds',
+    'get_round_factor',
 ]
 
 # Disable max row limitations in Altair.
@@ -139,12 +143,12 @@ class OptimizationGrid:
     return self._grid_dataset
 
   @property
-  def spend_grid(self) -> np.ndarray:
+  def spend_grid(self) -> xr.DataArray:
     """The spend grid."""
     return self.grid_dataset.spend_grid
 
   @property
-  def incremental_outcome_grid(self) -> np.ndarray:
+  def incremental_outcome_grid(self) -> xr.DataArray:
     """The incremental outcome grid."""
     return self.grid_dataset.incremental_outcome_grid
 
@@ -231,11 +235,7 @@ class OptimizationGrid:
             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)
+    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.'
@@ -244,7 +244,7 @@ class OptimizationGrid:
           ' It is only a problem when you use a much smaller budget, '
           ' for which the intended step size is smaller. '
       )
-    (spend_grid, incremental_outcome_grid) = self._trim_grid(
+    (spend_grid, incremental_outcome_grid) = self.trim_grids(
         spend_bound_lower=optimization_lower_bound,
         spend_bound_upper=optimization_upper_bound,
     )
@@ -267,86 +267,12 @@ class OptimizationGrid:
         },
     )
 
-  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` 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.
-    """
-    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(
-          spend=spend,
-          incremental_outcome=incremental_outcome,
-          roi_grid_point=roi_grid_point,
-          scenario=scenario,
-      ):
-        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 _trim_grid(
+  def trim_grids(
       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.
+  ) -> tuple[xr.DataArray, xr.DataArray]:
+    """Trims the grids based on a more restricted spend bound.
 
     Args:
       spend_bound_lower: The lower bound of spend for each channel. Must be in
@@ -355,12 +281,15 @@ class OptimizationGrid:
         the same order as `self.channels`.
 
     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.
+      updated_spend: The updated spend grid with only valid spend values.
+      updated_incremental_outcome: The updated incremental outcome grid
+        containing only the corresponding incremental outcome values for the
+        updated spend grid.
     """
+    self.check_optimization_bounds(
+        lower_bound=spend_bound_lower,
+        upper_bound=spend_bound_upper,
+    )
     spend_grid = self.spend_grid
     updated_spend = self.spend_grid.copy()
     updated_incremental_outcome = self.incremental_outcome_grid.copy()
@@ -387,6 +316,12 @@ class OptimizationGrid:
       updated_spend[nan_indices:, ch] = np.nan
       updated_incremental_outcome[nan_indices:, ch] = np.nan
 
+    # Drop the rows with all NaN values.
+    updated_spend = updated_spend.dropna(dim=c.GRID_SPEND_INDEX, how='all')
+    updated_incremental_outcome = updated_incremental_outcome.dropna(
+        dim=c.GRID_SPEND_INDEX, how='all'
+    )
+
     return (updated_spend, updated_incremental_outcome)
 
   def check_optimization_bounds(
@@ -429,6 +364,74 @@ class OptimizationGrid:
           + '\n'.join(errors)
       )
 
+  def _grid_search(
+      self,
+      spend_grid: xr.DataArray,
+      incremental_outcome_grid: xr.DataArray,
+      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:
+      `np.ndarray` of dimension (`n_total_channels`) containing the optimal
+      media spend that maximizes incremental outcome based on spend constraints
+      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(
+        backend.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(
+          spend=spend,
+          incremental_outcome=incremental_outcome,
+          roi_grid_point=roi_grid_point,
+          scenario=scenario,
+      ):
+        break
+
+      iterative_roi_grid[0 : row_idx + 1, media_idx] = np.nan
+      iterative_roi_grid[row_idx + 1 :, media_idx] = np.round(
+          backend.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:
@@ -881,9 +884,9 @@ class OptimizationResults:
     In particular:
 
     1. `spend_multiplier` matches the discrete optimization grid, considering
-       the grid step size and any channel-level constraint bounds.
+      the grid step size and any channel-level constraint bounds.
     2. `selected_times`, `by_reach`, and `use_optimal_frequency` match the
-       values set in `BudgetOptimizer.optimize()`.
+      values set in `BudgetOptimizer.optimize()`.
 
     Returns:
       A dataset returned by `Analyzer.response_curves()`, per budget
@@ -1329,25 +1332,25 @@ class BudgetOptimizer:
     The following optimization parameters are assigned default values based on
     the model input data:
     1. Flighting pattern. This is the relative allocation of a channel's media
-       units across geos and time periods. By default, the historical flighting
-       pattern is used. The default can be overridden by passing
-       `new_data.media`. The flighting pattern is held constant during
-       optimization and does not depend on the overall budget assigned to the
-       channel.
+      units across geos and time periods. By default, the historical flighting
+      pattern is used. The default can be overridden by passing
+      `new_data.media`. The flighting pattern is held constant during
+      optimization and does not depend on the overall budget assigned to the
+      channel.
     2. Cost per media unit. By default, the historical spend divided by
-       historical media units is used. This can optionally vary by geo or time
-       period or both depending on whether the spend data has geo and time
-       dimensions. The default can be overridden by passing `new_data.spend`.
-       The cost per media unit is held constant during optimization and does not
-       depend on the overall budget assigned to the channel.
+      historical media units is used. This can optionally vary by geo or time
+      period or both depending on whether the spend data has geo and time
+      dimensions. The default can be overridden by passing `new_data.spend`.
+      The cost per media unit is held constant during optimization and does not
+      depend on the overall budget assigned to the channel.
     3. Center of the spend box constraint for each channel. By default, the
-       historical percentage of spend within `selected_geos` and between
-       `start_date` and `end_date` is used. This can be overridden by passing
-       `pct_of_spend`.
+      historical percentage of spend within `selected_geos` and between
+      `start_date` and `end_date` is used. This can be overridden by passing
+      `pct_of_spend`.
     4. Total budget to be allocated (for fixed budget scenarios only). By
-       default, the historical spend within `selected_geos` and between
-       `start_date` and `end_date` is used. This can be overridden by passing
-       `budget`.
+      default, the historical spend within `selected_geos` and between
+      `start_date` and `end_date` is used. This can be overridden by passing
+      `budget`.
 
     Passing `new_data.media` (or `new_data.reach` or `new_data.frequency`) will
     override both the flighting pattern and cost per media unit. Passing
@@ -1530,7 +1533,8 @@ class BudgetOptimizer:
         use_kpi=use_kpi,
         hist_spend=optimization_grid.historical_spend,
         spend=spend.non_optimized,
-        selected_times=optimization_grid.selected_times,
+        start_date=start_date,
+        end_date=end_date,
         confidence_level=confidence_level,
         batch_size=batch_size,
         use_historical_budget=use_historical_budget,
@@ -1541,7 +1545,8 @@ class BudgetOptimizer:
         use_kpi=use_kpi,
         hist_spend=optimization_grid.historical_spend,
         spend=spend.non_optimized,
-        selected_times=optimization_grid.selected_times,
+        start_date=start_date,
+        end_date=end_date,
         optimal_frequency=optimization_grid.optimal_frequency,
         confidence_level=confidence_level,
         batch_size=batch_size,
@@ -1560,7 +1565,8 @@ class BudgetOptimizer:
         use_kpi=use_kpi,
         hist_spend=optimization_grid.historical_spend,
         spend=spend.optimized,
-        selected_times=optimization_grid.selected_times,
+        start_date=start_date,
+        end_date=end_date,
         optimal_frequency=optimization_grid.optimal_frequency,
         attrs=constraints,
         confidence_level=confidence_level,
@@ -1601,15 +1607,15 @@ class BudgetOptimizer:
 
   def create_optimization_tensors(
       self,
-      time: Sequence[str] | tf.Tensor,
-      cpmu: tf.Tensor | None = None,
-      media: tf.Tensor | None = None,
-      media_spend: tf.Tensor | None = None,
-      cprf: tf.Tensor | None = None,
-      rf_impressions: tf.Tensor | None = None,
-      frequency: tf.Tensor | None = None,
-      rf_spend: tf.Tensor | None = None,
-      revenue_per_kpi: tf.Tensor | None = None,
+      time: Sequence[str] | backend.Tensor,
+      cpmu: backend.Tensor | None = None,
+      media: backend.Tensor | None = None,
+      media_spend: backend.Tensor | None = None,
+      cprf: backend.Tensor | None = None,
+      rf_impressions: backend.Tensor | None = None,
+      frequency: backend.Tensor | None = None,
+      rf_spend: backend.Tensor | None = None,
+      revenue_per_kpi: backend.Tensor | None = None,
       use_optimal_frequency: bool = True,
   ) -> analyzer.DataTensors:
     """Creates a `DataTensors` for optimizations from CPM and flighting data.
@@ -1689,7 +1695,7 @@ class BudgetOptimizer:
         revenue_per_kpi=revenue_per_kpi,
         use_optimal_frequency=use_optimal_frequency,
     )
-    n_times = time.shape[0] if isinstance(time, tf.Tensor) else len(time)
+    n_times = time.shape[0] if isinstance(time, backend.Tensor) else len(time)
     n_geos = self._meridian.n_geos
     revenue_per_kpi = (
         _expand_tensor(revenue_per_kpi, (n_geos, n_times))
@@ -1714,25 +1720,25 @@ class BudgetOptimizer:
       )
       tensors[c.RF_SPEND] = allocated_impressions * cprf
       if use_optimal_frequency:
-        frequency = tf.ones_like(allocated_impressions)
+        frequency = backend.ones_like(allocated_impressions)
       tensors[c.FREQUENCY] = _expand_tensor(frequency, shape)
-      tensors[c.REACH] = tf.math.divide_no_nan(
+      tensors[c.REACH] = backend.divide_no_nan(
           allocated_impressions, tensors[c.FREQUENCY]
       )
     if rf_spend is not None:
       shape = (n_geos, n_times, rf_spend.shape[-1])
       cprf = _expand_tensor(cprf, shape)
       tensors[c.RF_SPEND] = self._allocate_tensor_by_population(rf_spend)
-      impressions = tf.math.divide_no_nan(tensors[c.RF_SPEND], cprf)
+      impressions = backend.divide_no_nan(tensors[c.RF_SPEND], cprf)
       if use_optimal_frequency:
-        frequency = tf.ones_like(impressions)
+        frequency = backend.ones_like(impressions)
       tensors[c.FREQUENCY] = _expand_tensor(frequency, shape)
-      tensors[c.REACH] = tf.math.divide_no_nan(
+      tensors[c.REACH] = backend.divide_no_nan(
           impressions, tensors[c.FREQUENCY]
       )
     if revenue_per_kpi is not None:
       tensors[c.REVENUE_PER_KPI] = revenue_per_kpi
-    tensors[c.TIME] = tf.convert_to_tensor(time)
+    tensors[c.TIME] = backend.to_tensor(time)
     return analyzer.DataTensors(**tensors)
 
   def _validate_grid(
@@ -1848,7 +1854,7 @@ class BudgetOptimizer:
       )
       return False
 
-    round_factor = _get_round_factor(budget, gtol)
+    round_factor = get_round_factor(budget, gtol)
     if round_factor != optimization_grid.round_factor:
       warnings.warn(
           'Optimization accuracy may suffer owing to budget level differences.'
@@ -1991,7 +1997,7 @@ class BudgetOptimizer:
         pct_of_spend=pct_of_spend,
     )
     spend = budget * valid_pct_of_spend
-    round_factor = _get_round_factor(budget, gtol)
+    round_factor = get_round_factor(budget, gtol)
     (optimization_lower_bound, optimization_upper_bound) = (
         get_optimization_bounds(
             n_channels=n_paid_channels,
@@ -2007,14 +2013,14 @@ class BudgetOptimizer:
           rf_spend=filled_data.rf_spend,
           revenue_per_kpi=filled_data.revenue_per_kpi,
       )
-      optimal_frequency = tf.convert_to_tensor(
+      optimal_frequency = backend.to_tensor(
           self._analyzer.optimal_freq(
               new_data=opt_freq_data,
               use_posterior=use_posterior,
               selected_times=selected_times,
               use_kpi=use_kpi,
           ).optimal_frequency,
-          dtype=tf.float32,
+          dtype=backend.float32,
       )
     else:
       optimal_frequency = None
@@ -2110,7 +2116,7 @@ class BudgetOptimizer:
       )
     else:
       assert new_data.time is not None
-      new_times_str = new_data.time.numpy().astype(str).tolist()
+      new_times_str = np.asarray(new_data.time).astype(str).tolist()
       time_coordinates = tc.TimeCoordinates.from_dates(new_times_str)
       expanded_dates = time_coordinates.expand_selected_time_dims(
           start_date=start_date,
@@ -2126,9 +2132,9 @@ class BudgetOptimizer:
       new_data: analyzer.DataTensors | None = None,
       optimal_frequency: Sequence[float] | None = None,
   ) -> tuple[
-      tf.Tensor | None,
-      tf.Tensor | None,
-      tf.Tensor | None,
+      backend.Tensor | None,
+      backend.Tensor | None,
+      backend.Tensor | None,
   ]:
     """Gets the tensors for incremental outcome, based on spend data.
 
@@ -2157,7 +2163,7 @@ class BudgetOptimizer:
         frequency is used for the optimization scenario.
 
     Returns:
-      Tuple of tf.tensors (new_media, new_reach, new_frequency).
+      Tuple of backend.tensors (new_media, new_reach, new_frequency).
     """
     new_data = new_data or analyzer.DataTensors()
     filled_data = new_data.validate_and_fill_missing_data(
@@ -2166,7 +2172,7 @@ class BudgetOptimizer:
     )
     if self._meridian.n_media_channels > 0:
       new_media = (
-          tf.math.divide_no_nan(
+          backend.divide_no_nan(
               spend[: self._meridian.n_media_channels],
               hist_spend[: self._meridian.n_media_channels],
           )
@@ -2177,7 +2183,7 @@ class BudgetOptimizer:
     if self._meridian.n_rf_channels > 0:
       rf_impressions = filled_data.reach * filled_data.frequency
       new_rf_impressions = (
-          tf.math.divide_no_nan(
+          backend.divide_no_nan(
               spend[-self._meridian.n_rf_channels :],
               hist_spend[-self._meridian.n_rf_channels :],
           )
@@ -2188,8 +2194,8 @@ class BudgetOptimizer:
           if optimal_frequency is None
           else optimal_frequency
       )
-      new_reach = tf.math.divide_no_nan(new_rf_impressions, frequency)
-      new_frequency = tf.math.divide_no_nan(new_rf_impressions, new_reach)
+      new_reach = backend.divide_no_nan(new_rf_impressions, frequency)
+      new_frequency = backend.divide_no_nan(new_rf_impressions, new_reach)
     else:
       new_reach = None
       new_frequency = None
@@ -2203,7 +2209,8 @@ class BudgetOptimizer:
       new_data: analyzer.DataTensors | None = None,
       use_posterior: bool = True,
       use_kpi: bool = False,
-      selected_times: Sequence[str] | Sequence[bool] | None = None,
+      start_date: tc.Date = None,
+      end_date: tc.Date = None,
       optimal_frequency: Sequence[float] | None = None,
       attrs: Mapping[str, Any] | None = None,
       confidence_level: float = c.DEFAULT_CONFIDENCE_LEVEL,
@@ -2216,8 +2223,11 @@ class BudgetOptimizer:
         c.PAID_DATA + (c.TIME,),
         self._meridian,
     )
-    spend_tensor = tf.convert_to_tensor(spend, dtype=tf.float32)
-    hist_spend = tf.convert_to_tensor(hist_spend, dtype=tf.float32)
+    selected_times = self._validate_selected_times(
+        start_date=start_date, end_date=end_date, new_data=new_data
+    )
+    spend_tensor = backend.to_tensor(spend, dtype=backend.float32)
+    hist_spend = backend.to_tensor(hist_spend, dtype=backend.float32)
     (new_media, new_reach, new_frequency) = (
         self._get_incremental_outcome_tensors(
             hist_spend,
@@ -2283,7 +2293,7 @@ class BudgetOptimizer:
     )
     effectiveness_with_mean_median_and_ci = (
         analyzer.get_central_tendency_and_ci(
-            data=tf.math.divide_no_nan(
+            data=backend.divide_no_nan(
                 incremental_outcome, aggregated_impressions
             ),
             confidence_level=confidence_level,
@@ -2292,12 +2302,12 @@ class BudgetOptimizer:
     )
 
     roi = analyzer.get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(incremental_outcome, spend_tensor),
+        data=backend.divide_no_nan(incremental_outcome, spend_tensor),
         confidence_level=confidence_level,
         include_median=True,
     )
     marginal_roi = analyzer.get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(
+        data=backend.divide_no_nan(
             mroi_numerator, spend_tensor * incremental_increase
         ),
         confidence_level=confidence_level,
@@ -2305,13 +2315,13 @@ class BudgetOptimizer:
     )
 
     cpik = analyzer.get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(spend_tensor, incremental_outcome),
+        data=backend.divide_no_nan(spend_tensor, incremental_outcome),
         confidence_level=confidence_level,
         include_median=True,
     )
-    total_inc_outcome = np.sum(incremental_outcome, -1)
-    total_cpik = np.mean(
-        tf.math.divide_no_nan(budget, total_inc_outcome),
+    total_inc_outcome = backend.reduce_sum(incremental_outcome, -1)
+    total_cpik = backend.reduce_mean(
+        backend.divide_no_nan(budget, total_inc_outcome),
         axis=(0, 1),
     )
 
@@ -2333,21 +2343,11 @@ class BudgetOptimizer:
         c.CPIK: ([c.CHANNEL, c.METRIC], cpik),
     }
 
-    all_times = (
-        filled_data.time.numpy().astype(str).tolist()
-        if filled_data.time is not None
-        else self._meridian.input_data.time.values.tolist()
-    )
-    if selected_times is not None and all(
-        isinstance(time, bool) for time in selected_times
-    ):
-      selected_times = [
-          time for time, selected in zip(all_times, selected_times) if selected
-      ]
+    all_times = np.asarray(filled_data.time).astype(str).tolist()
 
     attributes = {
-        c.START_DATE: min(selected_times) if selected_times else all_times[0],
-        c.END_DATE: max(selected_times) if selected_times else all_times[-1],
+        c.START_DATE: start_date if start_date else all_times[0],
+        c.END_DATE: end_date if end_date else all_times[-1],
         c.BUDGET: budget,
         c.PROFIT: total_incremental_outcome - budget,
         c.TOTAL_INCREMENTAL_OUTCOME: total_incremental_outcome,
@@ -2373,7 +2373,7 @@ class BudgetOptimizer:
       self,
       i: int,
       incremental_outcome_grid: np.ndarray,
-      multipliers_grid: tf.Tensor,
+      multipliers_grid: backend.Tensor,
       new_data: analyzer.DataTensors | None = None,
       selected_times: Sequence[str] | Sequence[bool] | None = None,
       use_posterior: bool = True,
@@ -2432,8 +2432,10 @@ class BudgetOptimizer:
       new_frequency = None
       new_reach = None
     elif optimal_frequency is not None:
-      new_frequency = tf.ones_like(filled_data.frequency) * optimal_frequency
-      new_reach = tf.math.divide_no_nan(
+      new_frequency = (
+          backend.ones_like(filled_data.frequency) * optimal_frequency
+      )
+      new_reach = backend.divide_no_nan(
           multipliers_grid[i, -self._meridian.n_rf_channels :]
           * filled_data.reach
           * filled_data.frequency,
@@ -2450,20 +2452,22 @@ class BudgetOptimizer:
     # (n_chains x n_draws x n_total_channels). Incremental_outcome_grid requires
     # incremental outcome by channel.
     incremental_outcome_grid[i, :] = np.mean(
-        self._analyzer.incremental_outcome(
-            use_posterior=use_posterior,
-            new_data=analyzer.DataTensors(
-                media=new_media,
-                reach=new_reach,
-                frequency=new_frequency,
-                revenue_per_kpi=filled_data.revenue_per_kpi,
-            ),
-            selected_times=selected_times,
-            use_kpi=use_kpi,
-            include_non_paid_channels=False,
-            batch_size=batch_size,
+        np.asarray(
+            self._analyzer.incremental_outcome(
+                use_posterior=use_posterior,
+                new_data=analyzer.DataTensors(
+                    media=new_media,
+                    reach=new_reach,
+                    frequency=new_frequency,
+                    revenue_per_kpi=filled_data.revenue_per_kpi,
+                ),
+                selected_times=selected_times,
+                use_kpi=use_kpi,
+                include_non_paid_channels=False,
+                batch_size=batch_size,
+            )
         ),
-        (c.CHAINS_DIMENSION, c.DRAWS_DIMENSION),
+        axis=(c.CHAINS_DIMENSION, c.DRAWS_DIMENSION),
         dtype=np.float64,
     )
 
@@ -2541,8 +2545,8 @@ class BudgetOptimizer:
       )
       spend_grid[: len(spend_grid_m), i] = spend_grid_m
     incremental_outcome_grid = np.full([n_grid_rows, n_grid_columns], np.nan)
-    multipliers_grid_base = tf.cast(
-        tf.math.divide_no_nan(spend_grid, spend), dtype=tf.float32
+    multipliers_grid_base = backend.cast(
+        backend.divide_no_nan(spend_grid, spend), dtype=backend.float32
     )
     multipliers_grid = np.where(
         np.isnan(spend_grid), np.nan, multipliers_grid_base
@@ -2573,7 +2577,7 @@ class BudgetOptimizer:
       rf_spend_max = np.nanmax(
           spend_grid[:, -self._meridian.n_rf_channels :], axis=0
       )
-      rf_roi = tf.math.divide_no_nan(rf_incremental_outcome_max, rf_spend_max)
+      rf_roi = backend.divide_no_nan(rf_incremental_outcome_max, rf_spend_max)
       incremental_outcome_grid[:, -self._meridian.n_rf_channels :] = (
           rf_roi * spend_grid[:, -self._meridian.n_rf_channels :]
       )
@@ -2581,14 +2585,14 @@ class BudgetOptimizer:
 
   def _validate_optimization_tensors(
       self,
-      cpmu: tf.Tensor | None = None,
-      cprf: tf.Tensor | None = None,
-      media: tf.Tensor | None = None,
-      rf_impressions: tf.Tensor | None = None,
-      frequency: tf.Tensor | None = None,
-      media_spend: tf.Tensor | None = None,
-      rf_spend: tf.Tensor | None = None,
-      revenue_per_kpi: tf.Tensor | None = None,
+      cpmu: backend.Tensor | None = None,
+      cprf: backend.Tensor | None = None,
+      media: backend.Tensor | None = None,
+      rf_impressions: backend.Tensor | None = None,
+      frequency: backend.Tensor | None = None,
+      media_spend: backend.Tensor | None = None,
+      rf_spend: backend.Tensor | None = None,
+      revenue_per_kpi: backend.Tensor | None = None,
       use_optimal_frequency: bool = True,
   ):
     """Validates the tensors needed for optimization."""
@@ -2642,7 +2646,7 @@ class BudgetOptimizer:
       )
 
   def _allocate_tensor_by_population(
-      self, tensor: tf.Tensor, required_ndim: int = 3
+      self, tensor: backend.Tensor, required_ndim: int = 3
   ):
     """Allocates a tensor of shape (time,) or (time, channel) by the population.
 
@@ -2664,16 +2668,83 @@ class BudgetOptimizer:
       )
 
     population = self._meridian.population
-    normalized_population = population / tf.reduce_sum(population)
+    normalized_population = population / backend.reduce_sum(population)
     if tensor.ndim == 1:
-      reshaped_population = normalized_population[:, tf.newaxis]
-      reshaped_tensor = tensor[tf.newaxis, :]
+      reshaped_population = normalized_population[:, backend.newaxis]
+      reshaped_tensor = tensor[backend.newaxis, :]
     else:
-      reshaped_population = normalized_population[:, tf.newaxis, tf.newaxis]
-      reshaped_tensor = tensor[tf.newaxis, :, :]
+      reshaped_population = normalized_population[
+          :, backend.newaxis, backend.newaxis
+      ]
+      reshaped_tensor = tensor[backend.newaxis, :, :]
     return reshaped_tensor * reshaped_population
 
 
+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 get_round_factor(budget: float, gtol: float) -> int:
+  """Gets the number of integer digits to round off of budget.
+
+  Args:
+    budget: Float number for total advertising budget.
+    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 int such that `(budget - rounded_budget) <= (budget * gtol)`.
+      `gtol` must be less than 1.
+
+  Returns:
+    Integer number of digits to round budget to.
+  """
+  tolerance = budget * gtol
+  if gtol >= 1.0:
+    raise ValueError('gtol must be less than one.')
+  elif budget <= 0.0:
+    raise ValueError('`budget` must be greater than zero.')
+  elif tolerance < 1.0:
+    return 0
+  else:
+    return -int(math.log10(tolerance)) - 1
+
+
 def _validate_pct_of_spend(
     n_channels: int,
     hist_spend: np.ndarray,
@@ -2748,7 +2819,7 @@ def _get_spend_bounds(
 
   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.
+      the untreated lower and upper bound spend for each media and RF channel.
   """
   (spend_const_lower, spend_const_upper) = _validate_spend_constraints(
       n_channels,
@@ -2762,47 +2833,6 @@ def _get_spend_bounds(
   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,
@@ -2836,30 +2866,6 @@ def _validate_budget(
       )
 
 
-def _get_round_factor(budget: float, gtol: float) -> int:
-  """Function for obtaining number of integer digits to round off of budget.
-
-  Args:
-    budget: float total advertising budget.
-    gtol: float indicating the acceptable relative error for the udget used in
-      the grid setup. The budget will be rounded by 10*n, where n is the
-      smallest int such that (budget - rounded_budget) is less than or equal to
-      (budget * gtol). gtol must be less than 1.
-
-  Returns:
-    int number of integer digits to round budget to.
-  """
-  tolerance = budget * gtol
-  if gtol >= 1.0:
-    raise ValueError('gtol must be less than one.')
-  elif budget <= 0.0:
-    raise ValueError('`budget` must be greater than zero.')
-  elif tolerance < 1.0:
-    return 0
-  else:
-    return -int(math.log10(tolerance)) - 1
-
-
 def _exceeds_optimization_constraints(
     spend: np.ndarray,
     incremental_outcome: np.ndarray,
@@ -2928,12 +2934,12 @@ def _raise_warning_if_target_constraints_not_met(
       )
 
 
-def _expand_tensor(tensor: tf.Tensor, required_shape: tuple[int, ...]):
+def _expand_tensor(tensor: backend.Tensor, required_shape: tuple[int, ...]):
   """Expands a tensor to the required number of dimensions."""
   if tensor.shape == required_shape:
     return tensor
   if tensor.ndim == 0:
-    return tf.fill(required_shape, tensor)
+    return backend.fill(required_shape, tensor)
 
   # Tensor must be less than or equal to the required number of dimensions and
   # the shape must match the required shape excluding the difference in number
@@ -2943,8 +2949,10 @@ def _expand_tensor(tensor: tf.Tensor, required_shape: tuple[int, ...]):
   ):
     n_tile_dims = len(required_shape) - tensor.ndim
     repeats = list(required_shape[:n_tile_dims]) + [1] * tensor.ndim
-    reshaped_tensor = tf.reshape(tensor, [1] * n_tile_dims + list(tensor.shape))
-    return tf.tile(reshaped_tensor, repeats)
+    reshaped_tensor = backend.reshape(
+        tensor, [1] * n_tile_dims + list(tensor.shape)
+    )
+    return backend.tile(reshaped_tensor, repeats)
 
   raise ValueError(
       f'Cannot expand tensor with shape {tensor.shape} to target'