Add ranked_geos property to InputData.

PiperOrigin-RevId: 822204516

Author: ez96

meridian/data/input_data.py

@@ -363,6 +363,45 @@ def media_time(self) -> xr.DataArray:
     else:
       return self.reach[constants.MEDIA_TIME]
 
+  @functools.cached_property
+  def ranked_geos(self) -> list[str]:
+    """Ranks geos by total spend then by total KPI."""
+    n_geos = len(self.geo)
+    if n_geos == 1:
+      return [self.geo[0]]
+    else:
+      if self.media_spend is not None and self.rf_spend is not None:
+        total_spend = (
+            self.media_spend.to_dataframe()
+            .unstack(level="media_channel")
+            .merge(
+                self.rf_spend.to_dataframe().unstack(level="rf_channel"),
+                on=[constants.GEO, constants.TIME],
+                how="inner",
+            )
+        )
+      elif self.media_spend is not None:
+        total_spend = self.media_spend.to_dataframe()
+      elif self.rf_spend is not None:
+        total_spend = self.rf_spend.to_dataframe()
+      else:
+        raise ValueError(
+            "It is required to have at least one of media or reach + frequency."
+        )
+
+      df_spend_sum = (
+          total_spend.sum(axis=1).rename("spend", inplace=True).to_frame()
+      )
+
+      return (
+          df_spend_sum.merge(
+              self.kpi.to_dataframe(), on=["geo", "time"], how="inner"
+          )
+          .groupby(constants.GEO)
+          .sum()
+          .sort_values(by=["spend", "kpi"], ascending=False)
+      ).index.values.tolist()
+
   @functools.cached_property
   def media_time_coordinates(self) -> tc.TimeCoordinates:
     """Returns the media time dimension in a `TimeCoordinates` wrapper."""

meridian/data/input_data_test.py

@@ -1642,6 +1642,82 @@ def test_scaled_centered_kpi_supports_dtype_int(self):
     data.population = data.population.astype(int)
     self.assertNotEmpty(data.scaled_centered_kpi)
 
+  def test_rank_geos_only_media_spend(self):
+    data = input_data.InputData(
+        kpi=self.not_lagged_kpi,
+        kpi_type=constants.NON_REVENUE,
+        population=self.population,
+        media=self.lagged_media,
+        media_spend=self.media_spend,
+    )
+    self.assertListEqual(
+        data.ranked_geos,
+        [
+            "geo_1",
+            "geo_5",
+            "geo_4",
+            "geo_6",
+            "geo_7",
+            "geo_8",
+            "geo_2",
+            "geo_3",
+            "geo_9",
+            "geo_0",
+        ],
+    )
+
+  def test_rank_geos_only_rf_spend(self):
+    data = input_data.InputData(
+        kpi=self.not_lagged_kpi,
+        kpi_type=constants.NON_REVENUE,
+        population=self.population,
+        reach=self.lagged_reach,
+        frequency=self.lagged_frequency,
+        rf_spend=self.rf_spend,
+    )
+    self.assertListEqual(
+        data.ranked_geos,
+        [
+            "geo_3",
+            "geo_5",
+            "geo_0",
+            "geo_4",
+            "geo_8",
+            "geo_6",
+            "geo_7",
+            "geo_1",
+            "geo_9",
+            "geo_2",
+        ],
+    )
+
+  def test_rank_geos_media_spend_and_rf_spend(self):
+    data = input_data.InputData(
+        kpi=self.not_lagged_kpi,
+        kpi_type=constants.NON_REVENUE,
+        population=self.population,
+        media=self.lagged_media,
+        media_spend=self.media_spend,
+        reach=self.lagged_reach,
+        frequency=self.lagged_frequency,
+        rf_spend=self.rf_spend,
+    )
+    self.assertListEqual(
+        data.ranked_geos,
+        [
+            "geo_5",
+            "geo_4",
+            "geo_1",
+            "geo_6",
+            "geo_8",
+            "geo_3",
+            "geo_7",
+            "geo_0",
+            "geo_2",
+            "geo_9",
+        ],
+    )
+
 
 class NonpaidInputDataTest(parameterized.TestCase):
   """Tests for non-paid InputData."""