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

meridian/analysis/analyzer.py

@@ -15,6 +15,7 @@
 """Methods to compute analysis metrics of the model and the data."""
 
 from collections.abc import Mapping, Sequence
+import dataclasses
 import itertools
 import numbers
 from typing import Any, Optional
@@ -53,6 +54,7 @@ def _validate_non_media_baseline_values_numbers(
 
 
 # TODO: Refactor the related unit tests to be under DataTensors.
+@dataclasses.dataclass
 class DataTensors(backend.ExtensionType):
   """Container for data variable arguments of Analyzer methods.
 
@@ -175,12 +177,31 @@ class DataTensors(backend.ExtensionType):
         else None
     )
     self.time = (
-        backend.to_tensor(time, dtype="string") if time is not None else None
+        backend.to_tensor(time, dtype=backend.string)
+        if time is not None
+        else None
     )
-
-  def __validate__(self):
     self._validate_n_dims()
 
+  def __eq__(self, other: Any) -> bool:
+    """Provides safe equality comparison for mixed tensor/non-tensor fields."""
+    if type(self) is not type(other):
+      return NotImplemented
+    for field in dataclasses.fields(self):
+      a = getattr(self, field.name)
+      b = getattr(other, field.name)
+      if a is None and b is None:
+        continue
+      if a is None or b is None:
+        return False
+      try:
+        if not bool(np.all(backend.to_tensor(backend.equal(a, b)))):
+          return False
+      except (ValueError, TypeError):
+        if a != b:
+          return False
+    return True
+
   def total_spend(self) -> backend.Tensor | None:
     """Returns the total spend tensor.
 
@@ -216,7 +237,7 @@ class DataTensors(backend.ExtensionType):
       of the corresponding tensor in the `meridian` object. If all time
       dimensions are the same, returns `None`.
     """
-    for field in self._tf_extension_type_fields():
+    for field in dataclasses.fields(self):
       new_tensor = getattr(self, field.name)
       if field.name == constants.RF_IMPRESSIONS:
         old_tensor = getattr(meridian.rf_tensors, field.name)
@@ -282,7 +303,7 @@ class DataTensors(backend.ExtensionType):
 
   def _validate_n_dims(self):
     """Raises an error if the tensors have the wrong number of dimensions."""
-    for field in self._tf_extension_type_fields():
+    for field in dataclasses.fields(self):
       tensor = getattr(self, field.name)
       if tensor is None:
         continue
@@ -315,7 +336,7 @@ class DataTensors(backend.ExtensionType):
       Warning: If an attribute exists in the `DataTensors` object that is not in
         the `required_variables` list, it will be ignored.
     """
-    for field in self._tf_extension_type_fields():
+    for field in dataclasses.fields(self):
       tensor = getattr(self, field.name)
       if tensor is None:
         continue
@@ -468,7 +489,7 @@ class DataTensors(backend.ExtensionType):
   ) -> Self:
     """Fills default values and returns a new DataTensors object."""
     output = {}
-    for field in self._tf_extension_type_fields():
+    for field in dataclasses.fields(self):
       var_name = field.name
       if var_name not in required_fields:
         continue
@@ -489,7 +510,7 @@ class DataTensors(backend.ExtensionType):
         old_tensor = meridian.revenue_per_kpi
       elif var_name == constants.TIME:
         old_tensor = backend.to_tensor(
-            meridian.input_data.time.values.tolist(), dtype="string"
+            meridian.input_data.time.values.tolist(), dtype=backend.string
         )
       else:
         continue
@@ -500,6 +521,7 @@ class DataTensors(backend.ExtensionType):
     return DataTensors(**output)
 
 
+@dataclasses.dataclass
 class DistributionTensors(backend.ExtensionType):
   """Container for parameters distributions arguments of Analyzer methods."""
 
@@ -583,17 +605,19 @@ def _transformed_new_or_scaled(
 
 def _calc_rsquared(expected, actual):
   """Calculates r-squared between actual and expected outcome."""
-  return 1 - np.nanmean((expected - actual) ** 2) / np.nanvar(actual)
+  return 1 - backend.nanmean((expected - actual) ** 2) / backend.nanvar(actual)
 
 
 def _calc_mape(expected, actual):
   """Calculates MAPE between actual and expected outcome."""
-  return np.nanmean(np.abs((actual - expected) / actual))
+  return backend.nanmean(backend.absolute((actual - expected) / actual))
 
 
 def _calc_weighted_mape(expected, actual):
   """Calculates wMAPE between actual and expected outcome (weighted by actual)."""
-  return np.nansum(np.abs(actual - expected)) / np.nansum(actual)
+  return backend.nansum(backend.absolute(actual - expected)) / backend.nansum(
+      actual
+  )
 
 
 def _warn_if_geo_arg_in_kwargs(**kwargs):
@@ -893,42 +917,37 @@ class Analyzer:
       )
     return result
 
-  def _check_revenue_data_exists(self, use_kpi: bool = False):
-    """Checks if the revenue data is available for the analysis.
+  def _use_kpi(self, use_kpi: bool = False) -> bool:
+    """Checks if KPI analysis should be used.
 
-    In the `kpi_type=NON_REVENUE` case, `revenue_per_kpi` is required to perform
-    the revenue analysis. If `revenue_per_kpi` is not defined, then the revenue
-    data is not available and the revenue analysis (`use_kpi=False`) is not
-    possible. Only the KPI analysis (`use_kpi=True`) is possible in this case.
+    If `use_kpi` is `True` but `kpi_type=REVENUE`, then `use_kpi` is ignored.
 
-    In the `kpi_type=REVENUE` case, KPI is equal to revenue and setting
-    `use_kpi=True` has no effect. Therefore, a warning is issued if the default
-    `False` value of `use_kpi` is overridden by the user.
+    If `use_kpi` is `False`, then  `revenue_per_kpi` is required to perform
+    the revenue analysis. Setting `use_kpi` to `False` in this case is ignored.
 
     Args:
-      use_kpi: A boolean flag indicating whether to use KPI instead of revenue.
+      use_kpi: A boolean flag indicating whether KPI analysis should be used.
 
+    Returns:
+      A boolean flag indicating whether KPI analysis should be used.
     Raises:
-      ValueError: If `use_kpi` is `False` and `revenue_per_kpi` is not defined.
-      UserWarning: If `use_kpi` is `True` in the `kpi_type=REVENUE` case.
+      UserWarning: If the KPI type is revenue and use_kpi is True or if
+      `use_kpi=False` but `revenue_per_kpi` is not available.
     """
-    if self._meridian.input_data.kpi_type == constants.NON_REVENUE:
-      if not use_kpi and self._meridian.revenue_per_kpi is None:
-        raise ValueError(
-            "Revenue analysis is not available when `revenue_per_kpi` is"
-            " unknown. Set `use_kpi=True` to perform KPI analysis instead."
-        )
+    if use_kpi and self._meridian.input_data.kpi_type == constants.REVENUE:
+      warnings.warn(
+          "Setting `use_kpi=True` has no effect when `kpi_type=REVENUE`"
+          " since in this case, KPI is equal to revenue."
+      )
+      return False
 
-    if self._meridian.input_data.kpi_type == constants.REVENUE:
-      # In the `kpi_type=REVENUE` case, KPI is equal to revenue and
-      # `revenue_per_kpi` is set to a tensor of 1s in the initialization of the
-      # `InputData` object.
-      assert self._meridian.revenue_per_kpi is not None
-      if use_kpi:
-        warnings.warn(
-            "Setting `use_kpi=True` has no effect when `kpi_type=REVENUE`"
-            " since in this case, KPI is equal to revenue."
-        )
+    if not use_kpi and self._meridian.input_data.revenue_per_kpi is None:
+      warnings.warn(
+          "Revenue analysis is not available when `revenue_per_kpi` is"
+          " unknown. Defaulting to KPI analysis."
+      )
+
+    return use_kpi or self._meridian.input_data.revenue_per_kpi is None
 
   def _get_adstock_dataframe(
       self,
@@ -1404,8 +1423,14 @@ class Analyzer:
             "`selected_geos` must match the geo dimension names from "
             "meridian.InputData."
         )
-      geo_mask = [x in selected_geos for x in mmm.input_data.geo]
-      tensor = backend.boolean_mask(tensor, geo_mask, axis=geo_dim)
+      geo_indices = [
+          i for i, x in enumerate(mmm.input_data.geo) if x in selected_geos
+      ]
+      tensor = backend.gather(
+          tensor,
+          backend.to_tensor(geo_indices, dtype=backend.int32),
+          axis=geo_dim,
+      )
 
     if selected_times is not None:
       _validate_selected_times(
@@ -1416,10 +1441,21 @@ class Analyzer:
           comparison_arg_name="`tensor`",
       )
       if _is_str_list(selected_times):
-        time_mask = [x in selected_times for x in mmm.input_data.time]
-        tensor = backend.boolean_mask(tensor, time_mask, axis=time_dim)
+        time_indices = [
+            i for i, x in enumerate(mmm.input_data.time) if x in selected_times
+        ]
+        tensor = backend.gather(
+            tensor,
+            backend.to_tensor(time_indices, dtype=backend.int32),
+            axis=time_dim,
+        )
       elif _is_bool_list(selected_times):
-        tensor = backend.boolean_mask(tensor, selected_times, axis=time_dim)
+        time_indices = [i for i, x in enumerate(selected_times) if x]
+        tensor = backend.gather(
+            tensor,
+            backend.to_tensor(time_indices, dtype=backend.int32),
+            axis=time_dim,
+        )
 
     tensor_dims = "...gt" + "m" * has_media_dim
     output_dims = (
@@ -1475,19 +1511,19 @@ class Analyzer:
         calculated.
       new_data: An optional `DataTensors` container with optional new tensors:
         `media`, `reach`, `frequency`, `organic_media`, `organic_reach`,
-        `organic_frequency`, `non_media_treatments`, `controls`. If `None`,
-        expected outcome is calculated conditional on the original values of the
-        data tensors that the Meridian object was initialized with. If
-        `new_data` argument is used, expected outcome is calculated conditional
-        on the values of the tensors passed in `new_data` and on the original
-        values of the remaining unset tensors. For example,
+        `organic_frequency`, `non_media_treatments`, `revenue_per_kpi`,
+        `controls`. If `None`, expected outcome is calculated conditional on the
+        original values of the data tensors that the Meridian object was
+        initialized with. If `new_data` argument is used, expected outcome is
+        calculated conditional on the values of the tensors passed in `new_data`
+        and on the original values of the remaining unset tensors. For example,
         `expected_outcome(new_data=DataTensors(reach=new_reach,
         frequency=new_frequency))` calculates expected outcome conditional on
         the original `media`, `organic_media`, `organic_reach`,
-        `organic_frequency`, `non_media_treatments` and `controls` tensors and
-        on the new given values for `reach` and `frequency` tensors. The new
-        tensors' dimensions must match the dimensions of the corresponding
-        original tensors from `input_data`.
+        `organic_frequency`, `non_media_treatments`, `revenue_per_kpi`, and
+        `controls` tensors and on the new given values for `reach` and
+        `frequency` tensors. The new tensors' dimensions must match the
+        dimensions of the corresponding original tensors from `input_data`.
       selected_geos: Optional list of containing a subset of geos to include. By
         default, all geos are included.
       selected_times: Optional list of containing a subset of dates to include.
@@ -1521,8 +1557,7 @@ class Analyzer:
         or `sample_prior()` (for `use_posterior=False`) has not been called
         prior to calling this method.
     """
-
-    self._check_revenue_data_exists(use_kpi)
+    use_kpi = self._use_kpi(use_kpi)
     self._check_kpi_transformation(inverse_transform_outcome, use_kpi)
     if self._meridian.is_national:
       _warn_if_geo_arg_in_kwargs(
@@ -1538,7 +1573,9 @@ class Analyzer:
     if new_data is None:
       new_data = DataTensors()
 
-    required_fields = constants.NON_REVENUE_DATA
+    required_fields = (
+        constants.PAID_DATA + constants.NON_PAID_DATA + (constants.CONTROLS,)
+    )
     filled_tensors = new_data.validate_and_fill_missing_data(
         required_tensors_names=required_fields,
         meridian=self._meridian,
@@ -1592,7 +1629,7 @@ class Analyzer:
     if inverse_transform_outcome:
       outcome_means = self._meridian.kpi_transformer.inverse(outcome_means)
       if not use_kpi:
-        outcome_means *= self._meridian.revenue_per_kpi
+        outcome_means *= filled_tensors.revenue_per_kpi
 
     return self.filter_and_aggregate_geos_and_times(
         outcome_means,
@@ -1721,7 +1758,7 @@ class Analyzer:
     Returns:
        Tensor of incremental outcome returned in terms of revenue or KPI.
     """
-    self._check_revenue_data_exists(use_kpi)
+    use_kpi = self._use_kpi(use_kpi)
     if revenue_per_kpi is None:
       revenue_per_kpi = self._meridian.revenue_per_kpi
     t1 = self._meridian.kpi_transformer.inverse(
@@ -1734,7 +1771,17 @@ class Analyzer:
       return kpi
     return backend.einsum("gt,...gtm->...gtm", revenue_per_kpi, kpi)
 
-  @backend.function(jit_compile=True)
+  @backend.function(
+      jit_compile=True,
+      static_argnames=[
+          "inverse_transform_outcome",
+          "use_kpi",
+          "selected_geos",
+          "selected_times",
+          "aggregate_geos",
+          "aggregate_times",
+      ],
+  )
   def _incremental_outcome_impl(
       self,
       data_tensors: DataTensors,
@@ -1804,7 +1851,7 @@ class Analyzer:
     Returns:
       Tensor containing the incremental outcome distribution.
     """
-    self._check_revenue_data_exists(use_kpi)
+    use_kpi = self._use_kpi(use_kpi)
     if (
         data_tensors.non_media_treatments is not None
         and non_media_treatments_baseline_normalized is None
@@ -2005,7 +2052,7 @@ class Analyzer:
         with matching time dimensions.
     """
     mmm = self._meridian
-    self._check_revenue_data_exists(use_kpi)
+    use_kpi = self._use_kpi(use_kpi)
     self._check_kpi_transformation(inverse_transform_outcome, use_kpi)
     if self._meridian.is_national:
       _warn_if_geo_arg_in_kwargs(
@@ -2146,8 +2193,12 @@ class Analyzer:
     )
     incremental_outcome_temps = [None] * len(batch_starting_indices)
     dim_kwargs = {
-        "selected_geos": selected_geos,
-        "selected_times": selected_times,
+        "selected_geos": (
+            tuple(selected_geos) if selected_geos is not None else None
+        ),
+        "selected_times": (
+            tuple(selected_times) if selected_times is not None else None
+        ),
         "aggregate_geos": aggregate_geos,
         "aggregate_times": aggregate_times,
     }
@@ -2322,7 +2373,7 @@ class Analyzer:
         "selected_times": selected_times,
         "aggregate_geos": aggregate_geos,
     }
-    self._check_revenue_data_exists(use_kpi)
+    use_kpi = self._use_kpi(use_kpi)
     self._validate_geo_and_time_granularity(**dim_kwargs)
     required_values = constants.PERFORMANCE_DATA
     if not new_data:
@@ -2431,6 +2482,7 @@ class Analyzer:
       (n_media_channels + n_rf_channels))`. The `n_geos` dimension is dropped if
       `aggregate_geos=True`.
     """
+    use_kpi = self._use_kpi(use_kpi)
     dim_kwargs = {
         "selected_geos": selected_geos,
         "selected_times": selected_times,
@@ -2444,7 +2496,6 @@ class Analyzer:
         "include_non_paid_channels": False,
         "aggregate_times": True,
     }
-    self._check_revenue_data_exists(use_kpi)
     self._validate_geo_and_time_granularity(**dim_kwargs)
     required_values = constants.PERFORMANCE_DATA
     if not new_data:
@@ -2632,6 +2683,7 @@ class Analyzer:
       self,
       aggregate_geos: bool = False,
       aggregate_times: bool = False,
+      use_kpi: bool = False,
       split_by_holdout_id: bool = False,
       non_media_baseline_values: Sequence[float] | None = None,
       confidence_level: float = constants.DEFAULT_CONFIDENCE_LEVEL,
@@ -2643,6 +2695,8 @@ class Analyzer:
         summed over all of the regions.
       aggregate_times: Boolean. If `True`, the expected, baseline, and actual
         are summed over all of the time periods.
+      use_kpi: If `True`, calculate the incremental KPI. Otherwise, calculate
+        the incremental revenue using the revenue per KPI (if available).
       split_by_holdout_id: Boolean. If `True` and `holdout_id` exists, the data
         is split into `'Train'`, `'Test'`, and `'All Data'` subsections.
       non_media_baseline_values: Optional list of shape
@@ -2659,8 +2713,8 @@ class Analyzer:
       A dataset with the expected, baseline, and actual outcome metrics.
     """
     _validate_non_media_baseline_values_numbers(non_media_baseline_values)
+    use_kpi = self._use_kpi(use_kpi)
     mmm = self._meridian
-    use_kpi = self._meridian.input_data.revenue_per_kpi is None
     can_split_by_holdout = self._can_split_by_holdout_id(split_by_holdout_id)
     expected_outcome = self.expected_outcome(
         aggregate_geos=False, aggregate_times=False, use_kpi=use_kpi
@@ -2828,7 +2882,7 @@ class Analyzer:
       self,
       use_posterior: bool,
       new_data: DataTensors | None = None,
-      use_kpi: bool | None = None,
+      use_kpi: bool = False,
       include_non_paid_channels: bool = True,
       non_media_baseline_values: Sequence[float] | None = None,
       **kwargs,
@@ -2875,7 +2929,7 @@ class Analyzer:
       the end containing the total incremental outcome of all channels.
     """
     _validate_non_media_baseline_values_numbers(non_media_baseline_values)
-    use_kpi = use_kpi or self._meridian.input_data.revenue_per_kpi is None
+    use_kpi = self._use_kpi(use_kpi)
     incremental_outcome_m = self.incremental_outcome(
         use_posterior=use_posterior,
         new_data=new_data,
@@ -3004,6 +3058,7 @@ class Analyzer:
       interpretation by time period.
     """
     _validate_non_media_baseline_values_numbers(non_media_baseline_values)
+    use_kpi = self._use_kpi(use_kpi)
     dim_kwargs = {
         "selected_geos": selected_geos,
         "selected_times": selected_times,
@@ -3146,16 +3201,19 @@ class Analyzer:
     ).where(lambda ds: ds.channel != constants.ALL_CHANNELS)
 
     if new_data.get_modified_times(self._meridian) is None:
+      expected_outcome_fields = list(
+          constants.PAID_DATA + constants.NON_PAID_DATA + (constants.CONTROLS,)
+      )
       expected_outcome_prior = self.expected_outcome(
           use_posterior=False,
-          new_data=new_data.filter_fields(constants.NON_REVENUE_DATA),
+          new_data=new_data.filter_fields(expected_outcome_fields),
           use_kpi=use_kpi,
           **dim_kwargs,
           **batched_kwargs,
       )
       expected_outcome_posterior = self.expected_outcome(
           use_posterior=True,
-          new_data=new_data.filter_fields(constants.NON_REVENUE_DATA),
+          new_data=new_data.filter_fields(expected_outcome_fields),
           use_kpi=use_kpi,
           **dim_kwargs,
           **batched_kwargs,
@@ -3399,6 +3457,7 @@ class Analyzer:
       aggregate_geos: bool = True,
       aggregate_times: bool = True,
       non_media_baseline_values: Sequence[float] | None = None,
+      use_kpi: bool = False,
       confidence_level: float = constants.DEFAULT_CONFIDENCE_LEVEL,
       batch_size: int = constants.DEFAULT_BATCH_SIZE,
   ) -> xr.Dataset:
@@ -3420,6 +3479,8 @@ class Analyzer:
         `model_spec.non_media_population_scaling_id` is `True`. If `None`, the
         `model_spec.non_media_baseline_values` is used, which defaults to the
         minimum value for each non_media treatment channel.
+      use_kpi: Boolean. If `True`, the baseline summary metrics are calculated
+        using KPI. If `False`, the metrics are calculated using revenue.
       confidence_level: Confidence level for media summary metrics credible
         intervals, represented as a value between zero and one.
       batch_size: Integer representing the maximum draws per chain in each
@@ -3435,7 +3496,7 @@ class Analyzer:
     _validate_non_media_baseline_values_numbers(non_media_baseline_values)
     # TODO: Change "pct_of_contribution" to a more accurate term.
 
-    use_kpi = self._meridian.input_data.revenue_per_kpi is None
+    use_kpi = self._use_kpi(use_kpi)
     dim_kwargs = {
         "selected_geos": selected_geos,
         "selected_times": selected_times,
@@ -3618,6 +3679,7 @@ class Analyzer:
       ValueError: If there are no channels with reach and frequency data.
     """
     dist_type = constants.POSTERIOR if use_posterior else constants.PRIOR
+    use_kpi = self._use_kpi(use_kpi)
     new_data = new_data or DataTensors()
     if self._meridian.n_rf_channels == 0:
       raise ValueError(
@@ -3696,9 +3758,11 @@ class Analyzer:
     )
 
     optimal_frequency = [freq_grid[i] for i in optimal_freq_idx]
-    optimal_frequency_tensor = backend.to_tensor(
-        backend.ones_like(filled_data.rf_impressions) * optimal_frequency,
-        backend.float32,
+    optimal_frequency_values = backend.to_tensor(
+        optimal_frequency, dtype=backend.float32
+    )
+    optimal_frequency_tensor = (
+        backend.ones_like(filled_data.rf_impressions) * optimal_frequency_values
     )
     optimal_reach = filled_data.rf_impressions / optimal_frequency_tensor
 
@@ -3783,10 +3847,7 @@ class Analyzer:
         attrs={
             constants.CONFIDENCE_LEVEL: confidence_level,
             constants.USE_POSTERIOR: use_posterior,
-            constants.IS_REVENUE_KPI: (
-                self._meridian.input_data.kpi_type == constants.REVENUE
-                or not use_kpi
-            ),
+            constants.IS_REVENUE_KPI: not use_kpi,
         },
     )
 
@@ -3794,6 +3855,7 @@ class Analyzer:
       self,
       selected_geos: Sequence[str] | None = None,
       selected_times: Sequence[str] | None = None,
+      use_kpi: bool = False,
       batch_size: int = constants.DEFAULT_BATCH_SIZE,
   ) -> xr.Dataset:
     """Calculates `R-Squared`, `MAPE`, and `wMAPE` goodness of fit metrics.
@@ -3824,6 +3886,8 @@ class Analyzer:
         default, all geos are included.
       selected_times: Optional list containing a subset of dates to include. By
         default, all time periods are included.
+      use_kpi: Whether to use KPI or revenue scale for the predictive accuracy
+        metrics.
       batch_size: Integer representing the maximum draws per chain in each
         batch. By default, `batch_size` is `100`. The calculation is run in
         batches to avoid memory exhaustion. If a memory error occurs, try
@@ -3837,7 +3901,7 @@ class Analyzer:
       is split into `'Train'`, `'Test'`, and `'All Data'` subsections, and the
       three metrics are computed for each.
     """
-    use_kpi = self._meridian.input_data.revenue_per_kpi is None
+    use_kpi = self._use_kpi(use_kpi)
     if self._meridian.is_national:
       _warn_if_geo_arg_in_kwargs(
           selected_geos=selected_geos,
@@ -3858,10 +3922,10 @@ class Analyzer:
         ],
         constants.GEO_GRANULARITY: [constants.GEO, constants.NATIONAL],
     }
-    if self._meridian.revenue_per_kpi is not None:
-      input_tensor = self._meridian.kpi * self._meridian.revenue_per_kpi
-    else:
+    if use_kpi:
       input_tensor = self._meridian.kpi
+    else:
+      input_tensor = self._meridian.kpi * self._meridian.revenue_per_kpi
     actual = np.asarray(
         self.filter_and_aggregate_geos_and_times(
             tensor=input_tensor,
@@ -3990,10 +4054,11 @@ class Analyzer:
           "sample_posterior() must be called prior to calling this method."
       )
 
-    def _transpose_first_two_dims(x: backend.Tensor) -> backend.Tensor:
-      n_dim = len(x.shape)
+    def _transpose_first_two_dims(x: Any) -> backend.Tensor:
+      x_tensor = backend.to_tensor(x)
+      n_dim = len(x_tensor.shape)
       perm = [1, 0] + list(range(2, n_dim))
-      return backend.transpose(x, perm)
+      return backend.transpose(x_tensor, perm)
 
     rhat = backend.mcmc.potential_scale_reduction({
         k: _transpose_first_two_dims(v)
@@ -4026,8 +4091,6 @@ class Analyzer:
     Returns:
       A DataFrame with the following columns:
 
-      *   `n_params`: The number of respective parameters in the model.
-      *   `avg_rhat`: The average R-hat value for the respective parameter.
       *   `n_params`: The number of respective parameters in the model.
       *   `avg_rhat`: The average R-hat value for the respective parameter.
       *   `max_rhat`: The maximum R-hat value for the respective parameter.
@@ -4944,11 +5007,12 @@ class Analyzer:
   def get_aggregated_spend(
       self,
       new_data: DataTensors | None = None,
+      selected_geos: Sequence[str] | None = None,
       selected_times: Sequence[str] | Sequence[bool] | None = None,
       include_media: bool = True,
       include_rf: bool = True,
   ) -> xr.DataArray:
-    """Gets the aggregated spend based on the selected time.
+    """Gets the aggregated spend based on the selected geos and time.
 
     Args:
       new_data: An optional `DataTensors` object containing the new `media`,
@@ -4959,6 +5023,9 @@ class Analyzer:
         of all the remaining tensors.  If any of the tensors in `new_data` is
         provided with a different number of time periods than in `InputData`,
         then all tensors must be provided with the same number of time periods.
+      selected_geos: Optional list containing a subset of geos to include. By
+        default, all geos are included. The selected geos should match those in
+        `InputData.geo`.
       selected_times: Optional list containing either a subset of dates to
         include or booleans with length equal to the number of time periods in
         KPI data. By default, all time periods are included.
@@ -5003,10 +5070,11 @@ class Analyzer:
       aggregated_media_spend = empty_da
     else:
       aggregated_media_spend = self._impute_and_aggregate_spend(
-          selected_times,
-          filled_data.media,
-          filled_data.media_spend,
-          list(self._meridian.input_data.media_channel.values),
+          selected_geos=selected_geos,
+          selected_times=selected_times,
+          media_execution_values=filled_data.media,
+          channel_spend=filled_data.media_spend,
+          channel_names=list(self._meridian.input_data.media_channel.values),
       )
 
     if not include_rf:
@@ -5025,10 +5093,11 @@ class Analyzer:
     else:
       rf_execution_values = filled_data.reach * filled_data.frequency
       aggregated_rf_spend = self._impute_and_aggregate_spend(
-          selected_times,
-          rf_execution_values,
-          filled_data.rf_spend,
-          list(self._meridian.input_data.rf_channel.values),
+          selected_geos=selected_geos,
+          selected_times=selected_times,
+          media_execution_values=rf_execution_values,
+          channel_spend=filled_data.rf_spend,
+          channel_names=list(self._meridian.input_data.rf_channel.values),
       )
 
     return xr.concat(
@@ -5037,21 +5106,26 @@ class Analyzer:
 
   def _impute_and_aggregate_spend(
       self,
+      selected_geos: Sequence[str] | None,
       selected_times: Sequence[str] | Sequence[bool] | None,
       media_execution_values: backend.Tensor,
       channel_spend: backend.Tensor,
       channel_names: Sequence[str],
   ) -> xr.DataArray:
-    """Imputes and aggregates the spend over the selected time period.
+    """Imputes and aggregates the spend within selected dimensions.
 
-    This function is used to aggregate the spend over the selected time period.
-    Imputation is required when `channel_spend` has only one dimension and the
-    aggregation is applied to only a subset of times, as specified by
-    `selected_times`. The `media_execution_values` argument only serves the
-    purpose of imputation. Although `media_execution_values` is a required
-    argument, its values only affect the output when imputation is required.
+    This function is used to aggregate the spend within selected geos over the
+    selected time period. Imputation is required when `channel_spend` has only
+    one dimension and the aggregation is applied to only a subset of geos or
+    times, as specified by `selected_geos` and `selected_times`. The
+    `media_execution_values` argument only serves the purpose of imputation.
+    Although `media_execution_values` is a required argument, its values only
+    affect the output when imputation is required.
 
     Args:
+      selected_geos: Optional list containing a subset of geos to include. By
+        default, all geos are included. The selected geos should match those in
+        `InputData.geo`.
       selected_times: Optional list containing either a subset of dates to
         include or booleans with length equal to the number of time periods in
         KPI data. By default, all time periods are included.
@@ -5066,7 +5140,7 @@ class Analyzer:
       variable `spend`.
     """
     dim_kwargs = {
-        "selected_geos": None,
+        "selected_geos": selected_geos,
         "selected_times": selected_times,
         "aggregate_geos": True,
         "aggregate_times": True,