google-meridian 1.1.6__py3-none-any.whl → 1.2.1__py3-none-any.whl

meridian/analysis/optimizer.py

@@ -24,15 +24,15 @@ 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 analyzer as analyzer_module
 from meridian.analysis import formatter
 from meridian.analysis import summary_text
 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.
@@ -110,7 +114,11 @@ class OptimizationGrid:
       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.
-    selected_times: The time coordinates from the model used in this grid.
+    selected_times: The time coordinates from the model used in this grid. If
+      new data with modified time coordinates is used for optimization, this
+      is a list of booleans indicating which time coordinates are selected.
+      Otherwise, this is a list of strings indicating the time coordinates used
+      in this grid.
   """
 
   _grid_dataset: xr.Dataset
@@ -124,7 +132,7 @@ class OptimizationGrid:
   gtol: float
   round_factor: int
   optimal_frequency: np.ndarray | None
-  selected_times: Sequence[str] | None
+  selected_times: Sequence[str] | Sequence[bool] | None
 
   @property
   def grid_dataset(self) -> xr.Dataset:
@@ -139,12 +147,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 +239,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 +248,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 +271,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 +285,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 +320,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,46 +368,107 @@ 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:
   """The optimized budget allocation.
 
   This is a dataclass object containing datasets output from `BudgetOptimizer`.
-  These datasets include:
-
-  - `nonoptimized_data`: The non-optimized budget metrics (based on historical
-    frequency).
-  - `nonoptimized_data_with_optimal_freq`: The non-optimized budget metrics
-    based on optimal frequency.
-  - `optimized_data`: The optimized budget metrics.
-  - `optimization_grid`: The grid information used for optimization.
 
-  The metrics (data variables) are: ROI, mROI, incremental outcome, CPIK.
-
-  Additionally, some intermediate values and referecences to the source fitted
-  model and analyzer are also stored here. These are useful for visualizing and
-  debugging.
+  The performance metrics (data variables) are: spend, percentage of spend, ROI,
+  mROI, incremental outcome, CPIK, and effectiveness.
 
   Attributes:
     meridian: The fitted Meridian model that was used to create this budget
       allocation.
     analyzer: The analyzer bound to the model above.
-    spend_ratio: The spend ratio used to scale the non-optimized budget metrics
-      to the optimized budget metrics.
-    spend_bounds: The spend bounds used to scale the non-optimized budget
-      metrics to the optimized budget metrics.
-    nonoptimized_data: The non-optimized budget metrics (based on historical
-      frequency).
-    nonoptimized_data_with_optimal_freq: The non-optimized budget metrics based
-      on optimal frequency.
-    optimized_data: The optimized budget metrics.
+    spend_ratio: The spend ratio used to scale the non-optimized performance
+      metrics to the optimized performance metrics.
+    spend_bounds: The spend bounds used to scale the non-optimized performance
+      metrics to the optimized performance metrics.
+    nonoptimized_data: Performance metrics under the non-optimized budget. For
+      R&F channels, the non-optimized frequency is used.
+    nonoptimized_data_with_optimal_freq: Performance metrics under the
+      non-optimized budget. For R&F channels, the optimal frequency is used if
+      frequency was optimized.
+    optimized_data: Performance metrics under the optimized budget. For R&F
+      channels, the optimal frequency is used if frequency was optimized.
     optimization_grid: The grid information used for optimization.
+    new_data: The optional `DataTensors` container that was used to create this
+      budget allocation.
   """
 
   meridian: model.Meridian
   # The analyzer bound to the model above.
-  analyzer: analyzer.Analyzer
+  analyzer: analyzer_module.Analyzer
   spend_ratio: np.ndarray  # spend / historical spend
   spend_bounds: tuple[np.ndarray, np.ndarray]
 
@@ -478,6 +478,10 @@ class OptimizationResults:
   _optimized_data: xr.Dataset
   _optimization_grid: OptimizationGrid
 
+  # The optional `DataTensors` container to use if optimization was performed
+  # on data different from the original `input_data`.
+  new_data: analyzer_module.DataTensors | None = None
+
   # TODO: Move this, and the plotting methods, to a summarizer.
   @functools.cached_property
   def template_env(self) -> jinja2.Environment:
@@ -494,10 +498,10 @@ class OptimizationResults:
 
   @property
   def nonoptimized_data(self) -> xr.Dataset:
-    """Dataset holding the non-optimized budget metrics.
+    """Dataset holding the non-optimized performance metrics.
 
     For channels that have reach and frequency data, their performance metrics
-    (ROI, mROI, incremental outcome, CPIK) are based on historical frequency.
+    are based on historical frequency.
 
     The dataset contains the following:
 
@@ -516,10 +520,10 @@ class OptimizationResults:
 
   @property
   def nonoptimized_data_with_optimal_freq(self) -> xr.Dataset:
-    """Dataset holding the non-optimized budget metrics.
+    """Dataset holding the non-optimized performance metrics.
 
     For channels that have reach and frequency data, their performance metrics
-    (ROI, mROI, incremental outcome, CPIK) are based on optimal frequency.
+    are based on optimal frequency.
 
     The dataset contains the following:
 
@@ -534,10 +538,10 @@ class OptimizationResults:
 
   @property
   def optimized_data(self) -> xr.Dataset:
-    """Dataset holding the optimized budget metrics.
+    """Dataset holding the optimized performance metrics.
 
     For channels that have reach and frequency data, their performance metrics
-    (ROI, mROI, incremental outcome) are based on optimal frequency.
+    are based on optimal frequency.
 
     The dataset contains the following:
 
@@ -881,9 +885,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
@@ -891,9 +895,12 @@ class OptimizationResults:
       returned this result.
     """
     channels = self.optimized_data.channel.values
-    selected_times = self.meridian.expand_selected_time_dims(
+    selected_times = _expand_selected_times(
+        meridian=self.meridian,
         start_date=self.optimized_data.start_date,
         end_date=self.optimized_data.end_date,
+        new_data=self.new_data,
+        return_flexible_str=True,
     )
     _, ubounds = self.spend_bounds
     upper_bound = (
@@ -909,6 +916,7 @@ class OptimizationResults:
     # WARN: If `selected_times` is not None (i.e. a subset time range), this
     # response curve computation might take a significant amount of time.
     return self.analyzer.response_curves(
+        new_data=self.new_data,
         spend_multipliers=spend_multiplier,
         use_posterior=self.optimization_grid.use_posterior,
         selected_times=selected_times,
@@ -1273,7 +1281,7 @@ class BudgetOptimizer:
 
   def __init__(self, meridian: model.Meridian):
     self._meridian = meridian
-    self._analyzer = analyzer.Analyzer(self._meridian)
+    self._analyzer = analyzer_module.Analyzer(self._meridian)
 
   def _validate_model_fit(self, use_posterior: bool):
     """Validates that the model is fit."""
@@ -1285,7 +1293,7 @@ class BudgetOptimizer:
 
   def optimize(
       self,
-      new_data: analyzer.DataTensors | None = None,
+      new_data: analyzer_module.DataTensors | None = None,
       use_posterior: bool = True,
       # TODO: b/409550413 - Remove this argument.
       selected_times: tuple[str | None, str | None] | None = None,
@@ -1329,25 +1337,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
@@ -1523,14 +1531,15 @@ class BudgetOptimizer:
     use_historical_budget = budget is None or np.isclose(
         budget, np.sum(optimization_grid.historical_spend)
     )
-    new_data = new_data or analyzer.DataTensors()
+    new_data = new_data or analyzer_module.DataTensors()
     nonoptimized_data = self._create_budget_dataset(
         new_data=new_data.filter_fields(c.PAID_DATA + (c.TIME,)),
         use_posterior=use_posterior,
         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 +1550,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 +1570,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,
@@ -1589,6 +1600,7 @@ class BudgetOptimizer:
     )
 
     return OptimizationResults(
+        new_data=new_data,
         meridian=self._meridian,
         analyzer=self._analyzer,
         spend_ratio=spend_ratio,
@@ -1601,17 +1613,17 @@ 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:
+  ) -> analyzer_module.DataTensors:
     """Creates a `DataTensors` for optimizations from CPM and flighting data.
 
     CPM is broken down into cost per media unit, `cpmu`, for the media channels
@@ -1689,7 +1701,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,30 +1726,30 @@ 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)
-    return analyzer.DataTensors(**tensors)
+    tensors[c.TIME] = backend.to_tensor(time)
+    return analyzer_module.DataTensors(**tensors)
 
   def _validate_grid(
       self,
-      new_data: analyzer.DataTensors | None,
+      new_data: analyzer_module.DataTensors | None,
       use_posterior: bool,
       start_date: tc.Date,
       end_date: tc.Date,
@@ -1793,7 +1805,7 @@ class BudgetOptimizer:
       return False
 
     if new_data is None:
-      new_data = analyzer.DataTensors()
+      new_data = analyzer_module.DataTensors()
     required_tensors = c.PERFORMANCE_DATA + (c.TIME,)
     filled_data = new_data.validate_and_fill_missing_data(
         required_tensors_names=required_tensors, meridian=self._meridian
@@ -1809,7 +1821,8 @@ class BudgetOptimizer:
       return False
 
     n_channels = len(optimization_grid.channels)
-    selected_times = self._validate_selected_times(
+    selected_times = _expand_selected_times(
+        meridian=self._meridian,
         start_date=start_date,
         end_date=end_date,
         new_data=new_data,
@@ -1848,7 +1861,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.'
@@ -1955,7 +1968,7 @@ class BudgetOptimizer:
     """
     self._validate_model_fit(use_posterior)
     if new_data is None:
-      new_data = analyzer.DataTensors()
+      new_data = analyzer_module.DataTensors()
 
     if selected_times is not None:
       warnings.warn(
@@ -1972,7 +1985,8 @@ class BudgetOptimizer:
     filled_data = new_data.validate_and_fill_missing_data(
         required_tensors_names=required_tensors, meridian=self._meridian
     )
-    selected_times = self._validate_selected_times(
+    selected_times = _expand_selected_times(
+        meridian=self._meridian,
         start_date=start_date,
         end_date=end_date,
         new_data=filled_data,
@@ -1991,7 +2005,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,
@@ -2002,19 +2016,19 @@ class BudgetOptimizer:
         )
     )
     if self._meridian.n_rf_channels > 0 and use_optimal_frequency:
-      opt_freq_data = analyzer.DataTensors(
+      opt_freq_data = analyzer_module.DataTensors(
           rf_impressions=filled_data.reach * filled_data.frequency,
           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
@@ -2092,43 +2106,16 @@ class BudgetOptimizer:
         attrs={c.SPEND_STEP_SIZE: spend_step_size},
     )
 
-  def _validate_selected_times(
-      self,
-      start_date: tc.Date,
-      end_date: tc.Date,
-      new_data: analyzer.DataTensors | None,
-  ) -> Sequence[str] | Sequence[bool] | None:
-    """Validates and returns the selected times."""
-    if start_date is None and end_date is None:
-      return None
-
-    new_data = new_data or analyzer.DataTensors()
-    if new_data.get_modified_times(self._meridian) is None:
-      return self._meridian.expand_selected_time_dims(
-          start_date=start_date,
-          end_date=end_date,
-      )
-    else:
-      assert new_data.time is not None
-      new_times_str = new_data.time.numpy().astype(str).tolist()
-      time_coordinates = tc.TimeCoordinates.from_dates(new_times_str)
-      expanded_dates = time_coordinates.expand_selected_time_dims(
-          start_date=start_date,
-          end_date=end_date,
-      )
-      expanded_str = [date.strftime(c.DATE_FORMAT) for date in expanded_dates]
-      return [x in expanded_str for x in new_times_str]
-
   def _get_incremental_outcome_tensors(
       self,
       hist_spend: np.ndarray,
       spend: np.ndarray,
-      new_data: analyzer.DataTensors | None = None,
+      new_data: analyzer_module.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,16 +2144,16 @@ 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()
+    new_data = new_data or analyzer_module.DataTensors()
     filled_data = new_data.validate_and_fill_missing_data(
         c.PAID_CHANNELS,
         self._meridian,
     )
     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 +2164,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 +2175,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
@@ -2200,10 +2187,11 @@ class BudgetOptimizer:
       self,
       hist_spend: np.ndarray,
       spend: np.ndarray,
-      new_data: analyzer.DataTensors | None = None,
+      new_data: analyzer_module.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,
@@ -2211,13 +2199,19 @@ class BudgetOptimizer:
       use_historical_budget: bool = True,
   ) -> xr.Dataset:
     """Creates the budget dataset."""
-    new_data = new_data or analyzer.DataTensors()
+    new_data = new_data or analyzer_module.DataTensors()
     filled_data = new_data.validate_and_fill_missing_data(
         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 = _expand_selected_times(
+        meridian=self._meridian,
+        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,
@@ -2227,7 +2221,7 @@ class BudgetOptimizer:
         )
     )
     budget = np.sum(spend_tensor)
-    inc_outcome_data = analyzer.DataTensors(
+    inc_outcome_data = analyzer_module.DataTensors(
         media=new_media,
         reach=new_reach,
         frequency=new_frequency,
@@ -2259,7 +2253,7 @@ class BudgetOptimizer:
     # shape (n_channels, n_metrics) where n_metrics = 4 for (mean, median,
     # ci_lo, and ci_hi)
     incremental_outcome_with_mean_median_and_ci = (
-        analyzer.get_central_tendency_and_ci(
+        analyzer_module.get_central_tendency_and_ci(
             data=incremental_outcome,
             confidence_level=confidence_level,
             include_median=True,
@@ -2271,7 +2265,7 @@ class BudgetOptimizer:
     )
 
     aggregated_impressions = self._analyzer.get_aggregated_impressions(
-        new_data=analyzer.DataTensors(
+        new_data=analyzer_module.DataTensors(
             media=new_media, reach=new_reach, frequency=new_frequency
         ),
         selected_times=selected_times,
@@ -2282,8 +2276,8 @@ class BudgetOptimizer:
         include_non_paid_channels=False,
     )
     effectiveness_with_mean_median_and_ci = (
-        analyzer.get_central_tendency_and_ci(
-            data=tf.math.divide_no_nan(
+        analyzer_module.get_central_tendency_and_ci(
+            data=backend.divide_no_nan(
                 incremental_outcome, aggregated_impressions
             ),
             confidence_level=confidence_level,
@@ -2291,27 +2285,27 @@ class BudgetOptimizer:
         )
     )
 
-    roi = analyzer.get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(incremental_outcome, spend_tensor),
+    roi = analyzer_module.get_central_tendency_and_ci(
+        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(
+    marginal_roi = analyzer_module.get_central_tendency_and_ci(
+        data=backend.divide_no_nan(
             mroi_numerator, spend_tensor * incremental_increase
         ),
         confidence_level=confidence_level,
         include_median=True,
     )
 
-    cpik = analyzer.get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(spend_tensor, incremental_outcome),
+    cpik = analyzer_module.get_central_tendency_and_ci(
+        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 +2327,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,8 +2357,8 @@ class BudgetOptimizer:
       self,
       i: int,
       incremental_outcome_grid: np.ndarray,
-      multipliers_grid: tf.Tensor,
-      new_data: analyzer.DataTensors | None = None,
+      multipliers_grid: backend.Tensor,
+      new_data: analyzer_module.DataTensors | None = None,
       selected_times: Sequence[str] | Sequence[bool] | None = None,
       use_posterior: bool = True,
       use_kpi: bool = False,
@@ -2416,7 +2400,7 @@ class BudgetOptimizer:
         reducing `batch_size`. The calculation will generally be faster with
         larger `batch_size` values.
     """
-    new_data = new_data or analyzer.DataTensors()
+    new_data = new_data or analyzer_module.DataTensors()
     filled_data = new_data.validate_and_fill_missing_data(
         c.PAID_DATA, self._meridian
     )
@@ -2432,8 +2416,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 +2436,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_module.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,
     )
 
@@ -2473,7 +2461,7 @@ class BudgetOptimizer:
       spend_bound_lower: np.ndarray,
       spend_bound_upper: np.ndarray,
       step_size: int,
-      new_data: analyzer.DataTensors | None = None,
+      new_data: analyzer_module.DataTensors | None = None,
       selected_times: Sequence[str] | Sequence[bool] | None = None,
       use_posterior: bool = True,
       use_kpi: bool = False,
@@ -2541,8 +2529,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 +2561,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 +2569,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 +2630,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 +2652,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 +2803,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 +2817,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 +2850,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 +2918,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,10 +2933,72 @@ 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'
       f' {required_shape}.'
   )
+
+
+def _expand_selected_times(
+    meridian: model.Meridian,
+    start_date: tc.Date,
+    end_date: tc.Date,
+    new_data: analyzer_module.DataTensors | None,
+    return_flexible_str: bool = False,
+) -> Sequence[str] | Sequence[bool] | None:
+  """Creates selected_times from start_date and end_date.
+
+  This function creates `selected_times` argument based on `start_date`,
+  `end_date` and `new_data`. If `new_data` is not used or used with unmodified
+  times, dates are selected from `meridian.input_data.time`. In the flexible
+  time scenario, when `new_data` is provided with modified times, dates are
+  selected from `new_data.time`. In this case, `new_data.time` must be provided
+  and the function returns a list of booleans.
+
+  Args:
+    meridian: The `Meridian` object with original data.
+    start_date: Start date of the selected time period.
+    end_date: End date of the selected time period.
+    new_data: The optional `DataTensors` object. If times are modified in
+      `new_data`, then `new_data.time` must be provided.
+    return_flexible_str: Whether to return a list of strings or a list of
+      booleans in case time is modified in `new_data`.
+
+  Returns:
+    If both `start_date` and `end_date` are `None`, returns `None`. If
+    `new_data` is not used or used with unmodified times, returns a list of
+    strings with selected dates. If `new_data` is used with modified times,
+    returns a list of strings or a list of booleans depending on the
+    `return_flexible_str` argument.
+  """
+  if start_date is None and end_date is None:
+    return None
+
+  new_data = new_data or analyzer_module.DataTensors()
+  if new_data.get_modified_times(meridian) is None:
+    return meridian.expand_selected_time_dims(
+        start_date=start_date,
+        end_date=end_date,
+    )
+  else:
+    assert new_data.time is not None
+    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,
+        end_date=end_date,
+    )
+    if return_flexible_str:
+      if expanded_dates is None:
+        expanded_dates = time_coordinates.all_dates
+      expanded_str = [date.strftime(c.DATE_FORMAT) for date in expanded_dates]
+      return [x for x in new_times_str if x in expanded_str]
+    # TODO: Remove once every method uses `new_data.time`.
+    else:
+      expanded_str = [date.strftime(c.DATE_FORMAT) for date in expanded_dates]
+      return [x in expanded_str for x in new_times_str]