Adds KPI invariance check to EDA engine.

PiperOrigin-RevId: 825871774

Author: riemanli

meridian/constants.py

@@ -133,6 +133,7 @@
 NON_MEDIA_TREATMENTS_SCALED = 'non_media_treatments_scaled'
 CONTROLS_SCALED = 'controls_scaled'
 KPI_SCALED = f'{KPI}_scaled'
+POPULATION_SCALED_KPI = f'{POPULATION}_scaled_{KPI}'
 RF_IMPRESSIONS_SCALED = f'{RF_IMPRESSIONS}_scaled'
 
 # Non-media treatments baseline value constants.

meridian/model/eda/__init__.py

@@ -15,4 +15,6 @@
 """The Meridian API module that performs EDA checks."""
 
 from meridian.model.eda import eda_engine
+from meridian.model.eda import eda_outcome
+from meridian.model.eda import eda_spec
 from meridian.model.eda import meridian_eda

meridian/model/eda/eda_engine.py

@@ -14,11 +14,15 @@
 
 """Meridian EDA Engine."""
 
+from __future__ import annotations
+
 import dataclasses
 import functools
+import typing
 from typing import Optional, Sequence
+
+from meridian import backend
 from meridian import constants
-from meridian.model import model
 from meridian.model import transformers
 from meridian.model.eda import eda_outcome
 from meridian.model.eda import eda_spec
@@ -30,6 +34,9 @@
 import xarray as xr
 
 
+if typing.TYPE_CHECKING:
+  from meridian.model import model  # pylint: disable=g-bad-import-order,g-import-not-at-top
+
 _DEFAULT_DA_VAR_AGG_FUNCTION = np.sum
 _CORRELATION_COL_NAME = 'correlation'
 _STACK_VAR_COORD_NAME = 'var'
@@ -134,7 +141,7 @@ class ReachFrequencyData:
 
 
 def _data_array_like(
-    *, da: xr.DataArray, values: np.ndarray | tf.Tensor
+    *, da: xr.DataArray, values: np.ndarray | backend.Tensor
 ) -> xr.DataArray:
   """Returns a DataArray from `values` with the same structure as `da`.
 
@@ -727,6 +734,27 @@ def geo_population_da(self) -> xr.DataArray | None:
         name=constants.POPULATION,
     )
 
+  @functools.cached_property
+  def _population_scaled_kpi_artifact(
+      self,
+  ) -> eda_outcome.KpiInvarianceArtifact:
+    """Returns an artifact containing population-scaled KPI data."""
+    kpi_transformer = self._meridian.kpi_transformer
+
+    population_scaled_kpi_da = _data_array_like(
+        da=self._meridian.input_data.kpi,
+        values=kpi_transformer.population_scaled_kpi,
+    )
+    population_scaled_kpi_da.name = constants.POPULATION_SCALED_KPI
+
+    artifact = eda_outcome.KpiInvarianceArtifact(
+        level=eda_outcome.AnalysisLevel.OVERALL,
+        population_scaled_kpi_da=population_scaled_kpi_da,
+        population_scaled_mean=float(kpi_transformer.population_scaled_mean),
+        population_scaled_stdev=float(kpi_transformer.population_scaled_stdev),
+    )
+    return artifact
+
   @functools.cached_property
   def kpi_scaled_da(self) -> xr.DataArray:
     scaled_kpi_da = _data_array_like(
@@ -1643,3 +1671,35 @@ def check_national_vif(
         findings=findings,
         analysis_artifacts=[national_vif_artifact],
     )
+
+  def kpi_has_variability(self) -> bool:
+    """Returns True if the KPI has variability across geos and times."""
+    stdev = float(self._meridian.kpi_transformer.population_scaled_stdev)
+    return stdev != 0
+
+  def check_overall_kpi_invariance(self) -> eda_outcome.EDAOutcome:
+    """Checks if the KPI is constant across all geos and times."""
+    if not self.kpi_has_variability():
+      kpi = 'kpi' if self._meridian.is_national else 'population_scaled_kpi'
+
+      eda_finding = eda_outcome.EDAFinding(
+          severity=eda_outcome.EDASeverity.ERROR,
+          explanation=(
+              f'`{kpi}` is constant across all geos and times, indicating no'
+              ' signal in the data. Please fix this data error.'
+          ),
+      )
+    else:
+      eda_finding = eda_outcome.EDAFinding(
+          severity=eda_outcome.EDASeverity.INFO,
+          explanation=(
+              'The KPI has variability across geos and times, indicating'
+              ' variability in the data.'
+          ),
+      )
+
+    return eda_outcome.EDAOutcome(
+        check_type=eda_outcome.EDACheckType.KPI_INVARIANCE,
+        findings=[eda_finding],
+        analysis_artifacts=[self._population_scaled_kpi_artifact],
+    )

meridian/model/eda/eda_engine_test.py

@@ -15,6 +15,7 @@
 from unittest import mock
 from absl.testing import absltest
 from absl.testing import parameterized
+from meridian import backend
 from meridian import constants
 from meridian.model import model
 from meridian.model import model_test_data
@@ -4605,6 +4606,113 @@ def test_check_national_vif_has_correct_vif_value(self):
     ]
     self.assertAllClose(national_artifact.vif_da.values, expected_national_vif)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="has_variability",
+          population_scaled_stdev=1.0,
+          expected_result=True,
+      ),
+      dict(
+          testcase_name="no_variability",
+          population_scaled_stdev=0.0,
+          expected_result=False,
+      ),
+  )
+  def test_kpi_has_variability(self, population_scaled_stdev, expected_result):
+    meridian = mock.Mock(spec=model.Meridian)
+    meridian.kpi_transformer.population_scaled_stdev = population_scaled_stdev
+    engine = eda_engine.EDAEngine(meridian)
+    self.assertEqual(engine.kpi_has_variability(), expected_result)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="geo",
+          is_national=False,
+          expected_kpi_name="population_scaled_kpi",
+      ),
+      dict(
+          testcase_name="national",
+          is_national=True,
+          expected_kpi_name="kpi",
+      ),
+  )
+  def test_check_overall_kpi_invariance_no_variability(
+      self, is_national, expected_kpi_name
+  ):
+    meridian = mock.Mock(spec=model.Meridian)
+    meridian.is_national = is_national
+    meridian.input_data.kpi = self.input_data_with_media_only.kpi
+    mock_kpi_transformer = mock.Mock()
+    mock_kpi_transformer.population_scaled_stdev = 0.0
+    mock_kpi_transformer.population_scaled_mean = 100.0
+    mock_kpi_transformer.population_scaled_kpi = backend.zeros(
+        (self._N_GEOS, self._N_TIMES), dtype=backend.float32
+    )
+    meridian.kpi_transformer = mock_kpi_transformer
+    engine = eda_engine.EDAEngine(meridian)
+
+    outcome = engine.check_overall_kpi_invariance()
+
+    self.assertEqual(
+        outcome.check_type, eda_outcome.EDACheckType.KPI_INVARIANCE
+    )
+    self.assertLen(outcome.findings, 1)
+    self.assertEqual(
+        outcome.findings[0].severity, eda_outcome.EDASeverity.ERROR
+    )
+    self.assertIn(
+        f"`{expected_kpi_name}` is constant across all geos and times",
+        outcome.findings[0].explanation,
+    )
+    self.assertLen(outcome.analysis_artifacts, 1)
+    artifact = outcome.analysis_artifacts[0]
+    self.assertIsInstance(artifact, eda_outcome.KpiInvarianceArtifact)
+    self.assertEqual(artifact.level, eda_outcome.AnalysisLevel.OVERALL)
+    self.assertEqual(artifact.population_scaled_stdev, 0.0)
+    self.assertEqual(artifact.population_scaled_mean, 100.0)
+    self.assertAllClose(
+        artifact.population_scaled_kpi_da.values,
+        backend.to_tensor(mock_kpi_transformer.population_scaled_kpi),
+    )
+
+  def test_check_overall_kpi_invariance_has_variability(self):
+    meridian = mock.Mock(spec=model.Meridian)
+    meridian.is_national = False
+    meridian.input_data.kpi = self.input_data_with_media_only.kpi
+    mock_kpi_transformer = mock.Mock()
+    mock_kpi_transformer.population_scaled_stdev = 1.0
+    mock_kpi_transformer.population_scaled_mean = 100.0
+    mock_kpi_transformer.population_scaled_kpi = backend.to_tensor(
+        np.arange(self._N_GEOS * self._N_TIMES).reshape(
+            self._N_GEOS, self._N_TIMES
+        ),
+        dtype=backend.float32,
+    )
+    meridian.kpi_transformer = mock_kpi_transformer
+    engine = eda_engine.EDAEngine(meridian)
+
+    outcome = engine.check_overall_kpi_invariance()
+
+    self.assertEqual(
+        outcome.check_type, eda_outcome.EDACheckType.KPI_INVARIANCE
+    )
+    self.assertLen(outcome.findings, 1)
+    self.assertEqual(outcome.findings[0].severity, eda_outcome.EDASeverity.INFO)
+    self.assertIn(
+        "The KPI has variability across geos and times",
+        outcome.findings[0].explanation,
+    )
+    self.assertLen(outcome.analysis_artifacts, 1)
+    artifact = outcome.analysis_artifacts[0]
+    self.assertIsInstance(artifact, eda_outcome.KpiInvarianceArtifact)
+    self.assertEqual(artifact.level, eda_outcome.AnalysisLevel.OVERALL)
+    self.assertEqual(artifact.population_scaled_stdev, 1.0)
+    self.assertEqual(artifact.population_scaled_mean, 100.0)
+    self.assertAllClose(
+        artifact.population_scaled_kpi_da.values,
+        backend.to_tensor(mock_kpi_transformer.population_scaled_kpi),
+    )
+
 
 if __name__ == "__main__":
   absltest.main()

meridian/model/eda/eda_outcome.py

@@ -122,13 +122,30 @@ class VIFArtifact(AnalysisArtifact):
   outlier_df: pd.DataFrame
 
 
+@dataclasses.dataclass(frozen=True)
+class KpiInvarianceArtifact(AnalysisArtifact):
+  """Encapsulates artifacts from a KPI invariance analysis.
+
+  Attributes:
+    population_scaled_kpi_da: DataArray of the population-scaled KPI.
+    population_scaled_mean: The mean of the population-scaled KPI.
+    population_scaled_stdev: The standard deviation of the population-scaled
+      KPI.
+  """
+
+  population_scaled_kpi_da: xr.DataArray
+  population_scaled_mean: float
+  population_scaled_stdev: float
+
+
 @enum.unique
 class EDACheckType(enum.Enum):
   """Enumeration for the type of an EDA check."""
 
   PAIRWISE_CORR = enum.auto()
   STD = enum.auto()
   VIF = enum.auto()
+  KPI_INVARIANCE = enum.auto()
 
 
 ArtifactType = typing.TypeVar('ArtifactType', bound='AnalysisArtifact')

meridian/model/eda/meridian_eda.py

@@ -13,12 +13,18 @@
 # limitations under the License.
 
 """Module containing Meridian related exploratory data analysis (EDA) functionalities."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 import altair as alt
-from meridian.model import model
+
+if TYPE_CHECKING:
+  from meridian.model import model
+
 
 __all__ = [
-    'MeridianEDA',
+    "MeridianEDA",
 ]
 
 

meridian/model/model.py

@@ -34,9 +34,10 @@
 from meridian.model import prior_sampler
 from meridian.model import spec
 from meridian.model import transformers
+from meridian.model.eda import eda_engine
+from meridian.model.eda import eda_spec as eda_spec_module
 import numpy as np
 
-
 __all__ = [
     "MCMCSamplingError",
     "MCMCOOMError",
@@ -91,6 +92,8 @@ class Meridian:
     model_spec: A `ModelSpec` object containing the model specification.
     inference_data: A _mutable_ `arviz.InferenceData` object containing the
       resulting data from fitting the model.
+    eda_engine: An `EDAEngine` object containing the EDA engine.
+    eda_spec: An `EDASpec` object containing the EDA specification.
     n_geos: Number of geos in the data.
     n_media_channels: Number of media channels in the data.
     n_rf_channels: Number of reach and frequency (RF) channels in the data.
@@ -154,12 +157,14 @@ def __init__(
       inference_data: (
           az.InferenceData | None
       ) = None,  # for deserializer use only
+      eda_spec: eda_spec_module.EDASpec = eda_spec_module.EDASpec(),
   ):
     self._input_data = input_data
     self._model_spec = model_spec if model_spec else spec.ModelSpec()
     self._inference_data = (
         inference_data if inference_data else az.InferenceData()
     )
+    self._eda_spec = eda_spec
 
     self._validate_data_dependent_model_spec()
     self._validate_injected_inference_data()
@@ -190,6 +195,14 @@ def model_spec(self) -> spec.ModelSpec:
   def inference_data(self) -> az.InferenceData:
     return self._inference_data
 
+  @functools.cached_property
+  def eda_engine(self) -> eda_engine.EDAEngine:
+    return eda_engine.EDAEngine(self, spec=self._eda_spec)
+
+  @property
+  def eda_spec(self) -> eda_spec_module.EDASpec:
+    return self._eda_spec
+
   @functools.cached_property
   def media_tensors(self) -> media.MediaTensors:
     return media.build_media_tensors(self.input_data, self.model_spec)
@@ -1142,13 +1155,9 @@ def _check_if_no_time_variation(
             " time."
         )
 
-  def _kpi_has_variability(self):
-    """Returns True if the KPI has variability across geos and times."""
-    return self.kpi_transformer.population_scaled_stdev != 0
-
   def _validate_kpi_transformer(self):
     """Validates the KPI transformer."""
-    if self._kpi_has_variability():
+    if self.eda_engine.kpi_has_variability():
       return
 
     kpi = "kpi" if self.is_national else "population_scaled_kpi"