google-meridian 1.0.4__py3-none-any.whl → 1.0.5__py3-none-any.whl

{google_meridian-1.0.4.dist-info → google_meridian-1.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: google-meridian
-Version: 1.0.4
+Version: 1.0.5
 Summary: Google's open source mixed marketing model library, helps you understand your return on investment and direct your ad spend with confidence. 
 Author-email: The Meridian Authors <no-reply@google.com>
 License: 
@@ -215,7 +215,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Topic :: Other/Nonlisted Topic
 Classifier: Topic :: Scientific/Engineering :: Information Analysis
-Requires-Python: >=3.11
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: arviz
@@ -392,7 +392,7 @@ To cite this repository:
   author = {Google Meridian Marketing Mix Modeling Team},
   title = {Meridian: Marketing Mix Modeling},
   url = {https://github.com/google/meridian},
-  version = {1.0.4},
+  version = {1.0.5},
   year = {2025},
 }
 ```

{google_meridian-1.0.4.dist-info → google_meridian-1.0.5.dist-info}/RECORD

@@ -1,9 +1,9 @@
-meridian/__init__.py,sha256=5Q57Y33UHL2lWhq6dNDShp8-jquMmYsIu4hiVfli9bg,714
+meridian/__init__.py,sha256=tjar6BXbxRQ0ihMDHu2kVYddlW0WGWsef0_5pw1VwOM,714
 meridian/constants.py,sha256=PhQX3S0b-Odv_bjl1YaO6vnLT4rBomiAcec0789VZNE,14599
 meridian/analysis/__init__.py,sha256=-FooDZ5OzePpyTVkvRoWQx_xBaRR_hjVLny9H8-kkyQ,836
-meridian/analysis/analyzer.py,sha256=8oLOroDg3Br8DQspiywYAVzCfidyU9a-HU-FVeKHpBc,204201
+meridian/analysis/analyzer.py,sha256=N-nWKmyJnXNv2RRt8bj-sO8hoBHUqMABU4Sj8eeBSNc,197121
 meridian/analysis/formatter.py,sha256=F8OYxD2bH13zV10JY63j2ugCOj-DpTXhyJr43n5ukr8,7270
-meridian/analysis/optimizer.py,sha256=GL65fy36x6DZqgepKrLjaia0o2pva6Um8nhXU_3MUVk,73058
+meridian/analysis/optimizer.py,sha256=Kbir76hEYAPvT5qstZ9edHwVkXs95KFJ6p66cCJqj-Y,77242
 meridian/analysis/summarizer.py,sha256=oJVqTCGLDf0cxh-1nFlAp4Hlo3NKXRYAoA67K9Hraw0,17498
 meridian/analysis/summary_text.py,sha256=wHdETMdnYiZn2wf5LfHzzcGbJMjTAI5GmYnketeJnU0,11469
 meridian/analysis/test_utils.py,sha256=WImpYOt5wLjhi8NczcxCvqSHRikTLvm1yli4o5DYzKE,77258
@@ -27,15 +27,15 @@ meridian/model/__init__.py,sha256=bvx8vvXolktsCTDKViU9U1v85pgNWF3haDowTKy11d4,98
 meridian/model/adstock_hill.py,sha256=b_YYhqci6ndgi602FFXmx2f12ceC4N0tp338nMMtm54,9283
 meridian/model/knots.py,sha256=r7PPaJM96d5pkoOeV9crIOgkM0-rh24mWMvypMiV4aQ,8054
 meridian/model/media.py,sha256=Gjr4jm0y_6pFy7aa_oKIuuZ8P7F56e3ZB-3o6msApeA,11876
-meridian/model/model.py,sha256=cs8YSq5kini5F_eQHqCj2eOYhgqMv3WndnzZEuNrw-U,43118
+meridian/model/model.py,sha256=WS1b_Gb-RlA3KPEaI6d1vsd78y09SYLEJhGpo4t0Mic,43178
 meridian/model/model_test_data.py,sha256=dqS_vDQUg811UGmyr8ZgWp8VTIra-krA7A2erQlfPlU,12488
 meridian/model/posterior_sampler.py,sha256=ffAzlkrCEqJjFmu6pDbUfam4BLTskxrSr0-U9-_aFiA,22637
 meridian/model/prior_distribution.py,sha256=6fqx_XIM0DSQICd65XaSRhelsjvZ4ariBfeyOeoKld8,39075
 meridian/model/prior_sampler.py,sha256=zGSAQviFO3s2GcVbfG9EfXxo_SNFBFbTQC3e-QBFzio,23079
 meridian/model/spec.py,sha256=xaHxfCLWLnWMAkMy2ouDoqGBHI_4tzzX8AaJOsKdu7Q,8878
 meridian/model/transformers.py,sha256=te3OJixprWLtv7O00a9GZWE4waTS94NNLVo3tWIl1-k,7420
-google_meridian-1.0.4.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-google_meridian-1.0.4.dist-info/METADATA,sha256=xkTYlkmsp-FG8Ovm9g1V2ZyoB5ncgPdNqmxSd6WRL6o,22040
-google_meridian-1.0.4.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-google_meridian-1.0.4.dist-info/top_level.txt,sha256=nwaCebZvvU34EopTKZsjK0OMTFjVnkf4FfnBN_TAc0g,9
-google_meridian-1.0.4.dist-info/RECORD,,
+google_meridian-1.0.5.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+google_meridian-1.0.5.dist-info/METADATA,sha256=gPEiheNFDpJRtx5a1tmjcgaXBYPh1ESH_gkTP5LTgLQ,22040
+google_meridian-1.0.5.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+google_meridian-1.0.5.dist-info/top_level.txt,sha256=nwaCebZvvU34EopTKZsjK0OMTFjVnkf4FfnBN_TAc0g,9
+google_meridian-1.0.5.dist-info/RECORD,,

meridian/__init__.py

@@ -14,7 +14,7 @@
 
 """Meridian API."""
 
-__version__ = "1.0.4"
+__version__ = "1.0.5"
 
 
 from meridian import analysis

meridian/analysis/analyzer.py

@@ -3480,175 +3480,6 @@ class Analyzer:
         baseline_pct_of_contribution,
     ])
 
-  # TODO: This method can be replaced once generalized
-  # `media_summary_metric` is done.
-  def _counterfactual_metric_dataset(
-      self,
-      use_posterior: bool = True,
-      new_data: DataTensors | None = None,
-      marginal_roi_by_reach: bool = True,
-      selected_geos: Sequence[str] | None = None,
-      selected_times: Sequence[str] | None = None,
-      use_kpi: bool = False,
-      attrs: Mapping[str, Any] | None = None,
-      confidence_level: float = constants.DEFAULT_CONFIDENCE_LEVEL,
-      batch_size: int = constants.DEFAULT_BATCH_SIZE,
-  ) -> xr.Dataset:
-    """Calculates the counterfactual metric dataset.
-
-    Args:
-      use_posterior: Boolean. If `True`, posterior counterfactual metrics are
-        generated. If `False`, prior counterfactual metrics are generated.
-      new_data: Optional DataTensors. When specified, it contains the
-        counterfactual `media`, `reach`, `frequency`, `media_spend`, `rf_spend`
-        and `revenue_per_kpi` values. The new tensors' dimensions must match the
-        dimensions of the corresponding original tensors from
-        `meridian.input_data`. If `None`, the existing tensors from the Meridian
-        object are used.
-      marginal_roi_by_reach: Boolean. Marginal ROI (mROI) is defined as the
-        return on the next dollar spent. If this argument is `True`, the
-        assumption is that the next dollar spent only impacts reach, holding
-        frequency constant. If this argument is `False`, the assumption is that
-        the next dollar spent only impacts frequency, holding reach constant.
-      selected_geos: Optional list contains a subset of geos to include. By
-        default, all geos are included.
-      selected_times: Optional list contains a subset of times to include. By
-        default, all time periods are included.
-      use_kpi: Boolean. If `True`, the counterfactual metrics are calculated
-        using KPI. If `False`, the counterfactual metrics are calculated using
-        revenue.
-      attrs: Optional dictionary of attributes to add to the dataset.
-      confidence_level: Confidence level for prior and posterior credible
-        intervals, represented as a value between zero and one.
-      batch_size: Maximum 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
-        larger `batch_size` values.
-
-    Returns:
-      An xarray Dataset which contains:
-      * Coordinates: `channel`, `metric` (`mean`, `median`, `ci_lo`, `ci_hi`).
-      * Data variables:
-        * `spend`: The spend for each channel.
-        * `pct_of_spend`: The percentage of spend for each channel.
-        * `incremental_outcome`: The incremental outcome for each channel.
-        * `pct_of_contribution`: The contribution percentage for each channel.
-        * `roi`: The ROI for each channel.
-        * `effectiveness`: The effectiveness for each channel.
-        * `mroi`: The marginal ROI for each channel.
-        * `cpik`: The CPIK for each channel.
-    """
-    dim_kwargs = {
-        "selected_geos": selected_geos,
-        "selected_times": selected_times,
-    }
-    metric_tensor_kwargs = {
-        "use_posterior": use_posterior,
-        "use_kpi": use_kpi,
-        "batch_size": batch_size,
-    }
-    filled_data = self._validate_and_fill_roi_analysis_arguments(
-        new_data=new_data or DataTensors()
-    )
-    spend = filled_data.total_spend()
-    if spend is not None and spend.ndim == 3:
-      spend = self.filter_and_aggregate_geos_and_times(spend, **dim_kwargs)
-
-    # _counterfactual_metric_dataset() is called only from `optimal_freq()`
-    # and uses only paid channels.
-    incremental_outcome_tensor = self.incremental_outcome(
-        new_data=filled_data,
-        include_non_paid_channels=False,
-        **dim_kwargs,
-        **metric_tensor_kwargs,
-    )
-    # expected_outcome returns a tensor of shape (n_chains, n_draws).
-    mean_expected_outcome = tf.reduce_mean(
-        self.expected_outcome(
-            new_data=filled_data,
-            **dim_kwargs,
-            **metric_tensor_kwargs,
-        ),
-        (0, 1),
-    )
-
-    # Calculate the mean, median, and confidence intervals for each metric.
-    incremental_outcome = get_central_tendency_and_ci(
-        data=incremental_outcome_tensor,
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-    pct_of_contribution = get_central_tendency_and_ci(
-        data=incremental_outcome_tensor
-        / mean_expected_outcome[..., None]
-        * 100,
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-    roi = get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(incremental_outcome_tensor, spend),
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-    mroi = get_central_tendency_and_ci(
-        data=self.marginal_roi(
-            by_reach=marginal_roi_by_reach,
-            new_data=filled_data,
-            **dim_kwargs,
-            **metric_tensor_kwargs,
-        ),
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-    effectiveness = get_central_tendency_and_ci(
-        data=incremental_outcome_tensor
-        / self.get_aggregated_impressions(
-            **dim_kwargs,
-            optimal_frequency=filled_data.frequency,
-            include_non_paid_channels=False,
-        ),
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-    cpik = get_central_tendency_and_ci(
-        data=tf.math.divide_no_nan(spend, incremental_outcome_tensor),
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-
-    budget = np.sum(spend) if np.sum(spend) > 0 else 1
-    dims = [constants.CHANNEL, constants.METRIC]
-    data_vars = {
-        constants.SPEND: ([constants.CHANNEL], spend),
-        constants.PCT_OF_SPEND: ([constants.CHANNEL], spend / budget),
-        constants.INCREMENTAL_OUTCOME: (dims, incremental_outcome),
-        constants.PCT_OF_CONTRIBUTION: (dims, pct_of_contribution),
-        constants.ROI: (dims, roi),
-        constants.MROI: (dims, mroi),
-        constants.EFFECTIVENESS: (dims, effectiveness),
-        constants.CPIK: (dims, cpik),
-    }
-
-    return xr.Dataset(
-        data_vars=data_vars,
-        coords={
-            constants.CHANNEL: (
-                [constants.CHANNEL],
-                self._meridian.input_data.get_all_paid_channels(),
-            ),
-            constants.METRIC: (
-                [constants.METRIC],
-                [
-                    constants.MEAN,
-                    constants.MEDIAN,
-                    constants.CI_LO,
-                    constants.CI_HI,
-                ],
-            ),
-        },
-        attrs=attrs,
-    )
-
   def optimal_freq(
       self,
       freq_grid: Sequence[float] | None = None,
@@ -3696,8 +3527,6 @@ class Analyzer:
         * `roi`: The ROI for each frequency value in `freq_grid`.
         * `optimized_incremental_outcome`: The incremental outcome based on the
             optimal frequency.
-        * `optimized_pct_of_contribution`: The contribution percentage based on
-            the optimal frequency.
         * `optimized_effectiveness`: The effectiveness based on the optimal
             frequency.
         * `optimized_roi`: The ROI based on the optimal frequency.
@@ -3770,8 +3599,7 @@ class Analyzer:
     )
 
     # Compute the optimized metrics based on the optimal frequency.
-    optimized_metrics_by_reach = self._counterfactual_metric_dataset(
-        use_posterior=use_posterior,
+    optimized_metrics_by_reach = self.summary_metrics(
         new_data=DataTensors(
             reach=optimal_reach, frequency=optimal_frequency_tensor
         ),
@@ -3779,9 +3607,11 @@ class Analyzer:
         selected_geos=selected_geos,
         selected_times=selected_times,
         use_kpi=use_kpi,
-    ).sel({constants.CHANNEL: rf_channel_values})
-    optimized_metrics_by_frequency = self._counterfactual_metric_dataset(
-        use_posterior=use_posterior,
+    ).sel({
+        constants.CHANNEL: rf_channel_values,
+        constants.DISTRIBUTION: dist_type,
+    })
+    optimized_metrics_by_frequency = self.summary_metrics(
         new_data=DataTensors(
             reach=optimal_reach, frequency=optimal_frequency_tensor
         ),
@@ -3789,7 +3619,10 @@ class Analyzer:
         selected_geos=selected_geos,
         selected_times=selected_times,
         use_kpi=use_kpi,
-    ).sel({constants.CHANNEL: rf_channel_values})
+    ).sel({
+        constants.CHANNEL: rf_channel_values,
+        constants.DISTRIBUTION: dist_type,
+    })
 
     data_vars = {
         constants.ROI: (
@@ -3804,10 +3637,6 @@ class Analyzer:
             [constants.RF_CHANNEL, constants.METRIC],
             optimized_metrics_by_reach.incremental_outcome.data,
         ),
-        constants.OPTIMIZED_PCT_OF_CONTRIBUTION: (
-            [constants.RF_CHANNEL, constants.METRIC],
-            optimized_metrics_by_reach.pct_of_contribution.data,
-        ),
         constants.OPTIMIZED_ROI: (
             (constants.RF_CHANNEL, constants.METRIC),
             optimized_metrics_by_reach.roi.data,

meridian/analysis/optimizer.py

@@ -45,6 +45,64 @@ alt.data_transformers.disable_max_rows()
 _SpendConstraint: TypeAlias = float | Sequence[float]
 
 
+@dataclasses.dataclass(frozen=True)
+class OptimizationGrid:
+  """Optimization grid information.
+
+  Attributes:
+    spend: ndarray of shape `(n_paid_channels,)` containing the spend allocation
+      for spend for all media and RF channels. The order matches
+      `InputData.get_all_paid_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.
+    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
+      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.
+  """
+
+  _grid_dataset: xr.Dataset
+
+  spend: np.ndarray
+  use_kpi: bool
+  use_posterior: bool
+  use_optimal_frequency: bool
+  round_factor: int
+  optimal_frequency: np.ndarray | None
+  selected_times: list[str] | None
+
+  @property
+  def grid_dataset(self) -> xr.Dataset:
+    """Dataset holding the grid information used for optimization.
+
+    The dataset contains the following:
+
+      - Coordinates:  `grid_spend_index`, `channel`
+      - Data variables: `spend_grid`, `incremental_outcome_grid`
+      - Attributes: `spend_step_size`
+    """
+    return self._grid_dataset
+
+  @property
+  def spend_grid(self) -> np.ndarray:
+    """The spend grid."""
+    return self.grid_dataset.spend_grid
+
+  @property
+  def incremental_outcome_grid(self) -> np.ndarray:
+    """The incremental outcome grid."""
+    return self.grid_dataset.incremental_outcome_grid
+
+  @property
+  def spend_step_size(self) -> float:
+    """The spend step size."""
+    return self.grid_dataset.attrs[c.SPEND_STEP_SIZE]
+
+
 @dataclasses.dataclass(frozen=True)
 class OptimizationResults:
   """The optimized budget allocation.
@@ -69,10 +127,6 @@ class OptimizationResults:
     meridian: The fitted Meridian model that was used to create this budget
       allocation.
     analyzer: The analyzer bound to the model above.
-    use_posterior: Whether the posterior distribution was used to optimize the
-      budget. If `False`, the prior distribution was used.
-    use_optimal_frequency: Whether optimal frequency was used to optimize the
-      budget.
     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
@@ -88,10 +142,6 @@ class OptimizationResults:
   meridian: model.Meridian
   # The analyzer bound to the model above.
   analyzer: analyzer.Analyzer
-
-  # The intermediate values used to derive the optimized budget allocation.
-  use_posterior: bool
-  use_optimal_frequency: bool
   spend_ratio: np.ndarray  # spend / historical spend
   spend_bounds: tuple[np.ndarray, np.ndarray]
 
@@ -99,7 +149,7 @@ class OptimizationResults:
   _nonoptimized_data: xr.Dataset
   _nonoptimized_data_with_optimal_freq: xr.Dataset
   _optimized_data: xr.Dataset
-  _optimization_grid: xr.Dataset
+  _optimization_grid: OptimizationGrid
 
   # TODO: Move this, and the plotting methods, to a summarizer.
   @functools.cached_property
@@ -174,15 +224,8 @@ class OptimizationResults:
     return self._optimized_data
 
   @property
-  def optimization_grid(self) -> xr.Dataset:
-    """Dataset holding the grid information used for optimization.
-
-    The dataset contains the following:
-
-      - Coordinates:  `grid_spend_index`, `channel`
-      - Data variables: `spend_grid`, `incremental_outcome_grid`
-      - Attributes: `spend_step_size`
-    """
+  def optimization_grid(self) -> OptimizationGrid:
+    """The grid information used for optimization."""
     return self._optimization_grid
 
   def output_optimization_summary(self, filename: str, filepath: str):
@@ -539,10 +582,10 @@ class OptimizationResults:
     # response curve computation might take a significant amount of time.
     return self.analyzer.response_curves(
         spend_multipliers=spend_multiplier,
-        use_posterior=self.use_posterior,
+        use_posterior=self.optimization_grid.use_posterior,
         selected_times=selected_times,
         by_reach=True,
-        use_optimal_frequency=self.use_optimal_frequency,
+        use_optimal_frequency=self.optimization_grid.use_optimal_frequency,
     )
 
   def _get_plottable_response_curves_df(
@@ -674,7 +717,6 @@ class OptimizationResults:
         id=summary_text.SCENARIO_PLAN_CARD_ID,
         title=summary_text.SCENARIO_PLAN_CARD_TITLE,
     )
-
     scenario_type = (
         summary_text.FIXED_BUDGET_LABEL.lower()
         if self.optimized_data.fixed_budget
@@ -891,6 +933,14 @@ class BudgetOptimizer:
     self._meridian = meridian
     self._analyzer = analyzer.Analyzer(self._meridian)
 
+  def _validate_model_fit(self, use_posterior: bool):
+    """Validates that the model is fit."""
+    dist_type = c.POSTERIOR if use_posterior else c.PRIOR
+    if dist_type not in self._meridian.inference_data.groups():
+      raise model.NotFittedModelError(
+          'Running budget optimization scenarios requires fitting the model.'
+      )
+
   def optimize(
       self,
       use_posterior: bool = True,
@@ -980,12 +1030,13 @@ class BudgetOptimizer:
       An `OptimizationResults` object containing optimized budget allocation
       datasets, along with some of the intermediate values used to derive them.
     """
-    dist_type = c.POSTERIOR if use_posterior else c.PRIOR
-    if dist_type not in self._meridian.inference_data.groups():
-      raise model.NotFittedModelError(
-          'Running budget optimization scenarios requires fitting the model.'
-      )
-    self._validate_budget(fixed_budget, budget, target_roi, target_mroi)
+    _validate_budget(
+        fixed_budget=fixed_budget,
+        budget=budget,
+        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(
@@ -994,28 +1045,17 @@ class BudgetOptimizer:
       )
     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)
-    step_size = 10 ** (-round_factor)
     rounded_spend = np.round(spend, round_factor).astype(int)
-    spend_ratio = np.divide(
-        spend,
-        hist_spend,
-        out=np.zeros_like(hist_spend, dtype=float),
-        where=hist_spend != 0,
-    )
     if self._meridian.n_rf_channels > 0 and use_optimal_frequency:
       optimal_frequency = tf.convert_to_tensor(
           self._analyzer.optimal_freq(
@@ -1037,34 +1077,30 @@ class BudgetOptimizer:
             fixed_budget=fixed_budget,
         )
     )
-    (spend_grid, incremental_outcome_grid) = self._create_grids(
+    optimization_grid = self.create_optimization_grid(
         spend=hist_spend,
         spend_bound_lower=optimization_lower_bound,
         spend_bound_upper=optimization_upper_bound,
-        step_size=step_size,
         selected_times=selected_time_dims,
+        round_factor=round_factor,
         use_posterior=use_posterior,
         use_kpi=use_kpi,
+        use_optimal_frequency=use_optimal_frequency,
         optimal_frequency=optimal_frequency,
         batch_size=batch_size,
     )
+    # TODO: b/375644691) - Move grid search to a OptimizationGrid class.
     optimal_spend = self._grid_search(
-        spend_grid=spend_grid,
-        incremental_outcome_grid=incremental_outcome_grid,
+        spend_grid=optimization_grid.spend_grid,
+        incremental_outcome_grid=optimization_grid.incremental_outcome_grid,
         budget=np.sum(rounded_spend),
         fixed_budget=fixed_budget,
         target_mroi=target_mroi,
         target_roi=target_roi,
     )
-
-    constraints = {
-        c.FIXED_BUDGET: fixed_budget,
-    }
-    if target_roi:
-      constraints[c.TARGET_ROI] = target_roi
-    elif target_mroi:
-      constraints[c.TARGET_MROI] = target_mroi
-
+    use_historical_budget = budget is None or round(budget) == round(
+        np.sum(hist_spend)
+    )
     nonoptimized_data = self._create_budget_dataset(
         use_posterior=use_posterior,
         use_kpi=use_kpi,
@@ -1086,6 +1122,13 @@ class BudgetOptimizer:
         batch_size=batch_size,
         use_historical_budget=use_historical_budget,
     )
+    constraints = {
+        c.FIXED_BUDGET: fixed_budget,
+    }
+    if target_roi:
+      constraints[c.TARGET_ROI] = target_roi
+    elif target_mroi:
+      constraints[c.TARGET_MROI] = target_mroi
     optimized_data = self._create_budget_dataset(
         use_posterior=use_posterior,
         use_kpi=use_kpi,
@@ -1098,18 +1141,16 @@ class BudgetOptimizer:
         batch_size=batch_size,
         use_historical_budget=use_historical_budget,
     )
-
-    optimization_grid = self._create_optimization_grid(
-        spend_grid=spend_grid,
-        spend_step_size=step_size,
-        incremental_outcome_grid=incremental_outcome_grid,
+    spend_ratio = np.divide(
+        spend,
+        hist_spend,
+        out=np.zeros_like(hist_spend, dtype=float),
+        where=hist_spend != 0,
     )
 
     return OptimizationResults(
         meridian=self._meridian,
         analyzer=self._analyzer,
-        use_posterior=use_posterior,
-        use_optimal_frequency=use_optimal_frequency,
         spend_ratio=spend_ratio,
         spend_bounds=spend_bounds,
         _nonoptimized_data=nonoptimized_data,
@@ -1118,7 +1159,83 @@ class BudgetOptimizer:
         _optimization_grid=optimization_grid,
     )
 
-  def _create_optimization_grid(
+  def create_optimization_grid(
+      self,
+      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,
+      use_kpi: bool = False,
+      use_optimal_frequency: bool = True,
+      optimal_frequency: xr.DataArray | None = None,
+      batch_size: int = c.DEFAULT_BATCH_SIZE,
+  ) -> OptimizationGrid:
+    """Creates a OptimizationGrid for optimization.
+
+    Args:
+      spend: ndarray of shape `(n_paid_channels,)` with 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.
+      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
+        larger `batch_size` values.
+
+    Returns:
+      An OptimizationGrid object containing the grid data for optimization.
+    """
+    self._validate_model_fit(use_posterior)
+
+    step_size = 10 ** (-round_factor)
+    (spend_grid, incremental_outcome_grid) = self._create_grids(
+        spend=spend,
+        spend_bound_lower=spend_bound_lower,
+        spend_bound_upper=spend_bound_upper,
+        step_size=step_size,
+        selected_times=selected_times,
+        use_posterior=use_posterior,
+        use_kpi=use_kpi,
+        optimal_frequency=optimal_frequency,
+        batch_size=batch_size,
+    )
+    grid_dataset = self._create_grid_dataset(
+        spend_grid=spend_grid,
+        spend_step_size=step_size,
+        incremental_outcome_grid=incremental_outcome_grid,
+    )
+
+    return OptimizationGrid(
+        _grid_dataset=grid_dataset,
+        spend=spend,
+        use_kpi=use_kpi,
+        use_posterior=use_posterior,
+        use_optimal_frequency=use_optimal_frequency,
+        round_factor=round_factor,
+        optimal_frequency=optimal_frequency,
+        selected_times=selected_times,
+    )
+
+  def _create_grid_dataset(
       self,
       spend_grid: np.ndarray,
       spend_step_size: float,
@@ -1164,39 +1281,6 @@ class BudgetOptimizer:
         attrs={c.SPEND_STEP_SIZE: spend_step_size},
     )
 
-  def _validate_budget(
-      self,
-      fixed_budget: bool,
-      budget: float | None,
-      target_roi: float | None,
-      target_mroi: float | None,
-  ):
-    """Validates the budget optimization arguments."""
-    if fixed_budget:
-      if target_roi is not None:
-        raise ValueError(
-            '`target_roi` is only used for flexible budget scenarios.'
-        )
-      if target_mroi is not None:
-        raise ValueError(
-            '`target_mroi` is only used for flexible budget scenarios.'
-        )
-      if budget is not None and budget <= 0:
-        raise ValueError('`budget` must be greater than zero.')
-    else:
-      if budget is not None:
-        raise ValueError('`budget` is only used for fixed budget scenarios.')
-      if target_roi is None and target_mroi is None:
-        raise ValueError(
-            'Must specify either `target_roi` or `target_mroi` for flexible'
-            ' budget optimization.'
-        )
-      if target_roi is not None and target_mroi is not None:
-        raise ValueError(
-            'Must specify only one of `target_roi` or `target_mroi` for'
-            'flexible budget optimization.'
-        )
-
   def _validate_pct_of_spend(
       self, hist_spend: np.ndarray, pct_of_spend: Sequence[float] | None
   ) -> np.ndarray:
@@ -1390,27 +1474,6 @@ class BudgetOptimizer:
         incremental_outcome_with_mean_median_and_ci[:, 0]
     )
 
-    # expected_outcome here is a tensor with the shape (n_chains, n_draws)
-    expected_outcome = self._analyzer.expected_outcome(
-        use_posterior=use_posterior,
-        new_data=analyzer.DataTensors(
-            media=new_media,
-            reach=new_reach,
-            frequency=new_frequency,
-        ),
-        selected_times=selected_times,
-        use_kpi=use_kpi,
-        batch_size=batch_size,
-    )
-    mean_expected_outcome = tf.reduce_mean(expected_outcome, (0, 1))  # a scalar
-
-    pct_contrib = incremental_outcome / mean_expected_outcome[..., None] * 100
-    pct_contrib_with_mean_median_and_ci = analyzer.get_central_tendency_and_ci(
-        data=pct_contrib,
-        confidence_level=confidence_level,
-        include_median=True,
-    )
-
     aggregated_impressions = self._analyzer.get_aggregated_impressions(
         selected_times=selected_times,
         selected_geos=None,
@@ -1471,10 +1534,6 @@ class BudgetOptimizer:
             [c.CHANNEL, c.METRIC],
             incremental_outcome_with_mean_median_and_ci,
         ),
-        c.PCT_OF_CONTRIBUTION: (
-            [c.CHANNEL, c.METRIC],
-            pct_contrib_with_mean_median_and_ci,
-        ),
         c.EFFECTIVENESS: (
             [c.CHANNEL, c.METRIC],
             effectiveness_with_mean_median_and_ci,
@@ -1714,9 +1773,12 @@ 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 = tf.cast(
+    multipliers_grid_base = tf.cast(
         tf.math.divide_no_nan(spend_grid, spend), dtype=tf.float32
     )
+    multipliers_grid = np.where(
+        np.isnan(spend_grid), np.nan, multipliers_grid_base
+    )
     for i in range(n_grid_rows):
       self._update_incremental_outcome_grid(
           i=i,
@@ -1838,6 +1900,39 @@ class BudgetOptimizer:
     return spend_optimal
 
 
+def _validate_budget(
+    fixed_budget: bool,
+    budget: float | None,
+    target_roi: float | None,
+    target_mroi: float | None,
+):
+  """Validates the budget optimization arguments."""
+  if fixed_budget:
+    if target_roi is not None:
+      raise ValueError(
+          '`target_roi` is only used for flexible budget scenarios.'
+      )
+    if target_mroi is not None:
+      raise ValueError(
+          '`target_mroi` is only used for flexible budget scenarios.'
+      )
+    if budget is not None and budget <= 0:
+      raise ValueError('`budget` must be greater than zero.')
+  else:
+    if budget is not None:
+      raise ValueError('`budget` is only used for fixed budget scenarios.')
+    if target_roi is None and target_mroi is None:
+      raise ValueError(
+          'Must specify either `target_roi` or `target_mroi` for flexible'
+          ' budget optimization.'
+      )
+    if target_roi is not None and target_mroi is not None:
+      raise ValueError(
+          'Must specify only one of `target_roi` or `target_mroi` for'
+          'flexible budget optimization.'
+      )
+
+
 def _get_round_factor(budget: float, gtol: float) -> int:
   """Function for obtaining number of integer digits to round off of budget.
 
@@ -1902,6 +1997,11 @@ def _exceeds_optimization_constraints(
   if fixed_budget:
     return np.sum(spend) > budget
   elif target_roi is not None:
-    return (np.sum(incremental_outcome) / np.sum(spend)) < target_roi
+    cur_total_roi = np.sum(incremental_outcome) / np.sum(spend)
+    # In addition to the total roi being less than the target roi, the roi of
+    # the current optimization step should also be less than the total roi.
+    # Without the second condition, the optimization algorithm may not have
+    # found the roi point close to the target roi yet.
+    return cur_total_roi < target_roi and roi_grid_point < cur_total_roi
   else:
     return roi_grid_point < target_mroi

meridian/model/model.py

@@ -537,9 +537,10 @@ class Meridian:
     self._validate_injected_inference_data_group_coord(
         inference_data, group, constants.TIME, self.n_times
     )
-    self._validate_injected_inference_data_group_coord(
-        inference_data, group, constants.SIGMA_DIM, self._sigma_shape
-    )
+    if not self.model_spec.unique_sigma_for_each_geo:
+      self._validate_injected_inference_data_group_coord(
+          inference_data, group, constants.SIGMA_DIM, self._sigma_shape
+      )
     self._validate_injected_inference_data_group_coord(
         inference_data,
         group,