Add ranked_geos property to InputData.

PiperOrigin-RevId: 822204516

Author: ez96

meridian/data/input_data.py

@@ -363,6 +363,66 @@ 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:
+      spend_da = None
+
+      if self.media_spend is not None:
+        # Sum across media_channel and any other non-geo/time dimensions
+        media_spend_sum = self.media_spend.sum(
+            dim=[
+                d
+                for d in self.media_spend.dims
+                if d not in (constants.GEO, constants.TIME)
+            ]
+        )
+        spend_da = media_spend_sum
+
+      if self.rf_spend is not None:
+        # Sum across rf_channel and any other non-geo/time dimensions
+        rf_spend_sum = self.rf_spend.sum(
+            dim=[
+                d
+                for d in self.rf_spend.dims
+                if d not in (constants.GEO, constants.TIME)
+            ]
+        )
+
+        if spend_da is None:
+          spend_da = rf_spend_sum
+        else:
+          # Add media and rf spend. xarray aligns on geo/time.
+          # Use fillna(0) for safe addition.
+          spend_da = spend_da.fillna(0) + rf_spend_sum.fillna(0)
+
+      if spend_da is None:
+        raise ValueError(
+            "It is required to have at least one of media spend or rf spend."
+        )
+
+      # 2. Calculate Total Spend and KPI per Geo
+      geo_spend_sum_da = spend_da.sum(dim=constants.TIME)
+      geo_kpi_sum_da = self.kpi.sum(dim=constants.TIME)
+
+      # 3. Get the underlying NumPy arrays and coordinates
+      geo_coords = geo_spend_sum_da[constants.GEO].values
+      spend_values = geo_spend_sum_da.values
+      kpi_values = geo_kpi_sum_da.values
+
+      # 4. Use NumPy's lexsort for multi-criteria sorting (descending)
+      sort_indices = np.lexsort((
+          -kpi_values,  # Secondary sort key (minor)
+          -spend_values,  # Primary sort key (major)
+      ))
+
+      # 5. Apply the sorted indices to the geo coordinates
+      return geo_coords[sort_indices].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,96 @@ 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",
+        ],
+    )
+
+  def test_rank_geos_national(self):
+    data = test_utils.sample_input_data_from_dataset(
+        test_utils.random_dataset(
+            n_geos=1,
+            n_times=20,
+            n_media_times=20,
+            n_controls=2,
+            n_media_channels=5,
+        ),
+        "non_revenue",
+    )
+    actual_ranked = data.ranked_geos
+    self.assertListEqual(actual_ranked, ["geo_0"])
+
 
 class NonpaidInputDataTest(parameterized.TestCase):
   """Tests for non-paid InputData."""