Add data-to-parameter ratio check to EDA engine.

PiperOrigin-RevId: 893707590

Author: riemanli

CHANGELOG.md

@@ -23,6 +23,7 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+* Add EDA check for data-to-parameter ratio (DATA_ADEQUACY).
 * Add JAX 64-bit precision opt-in configuration.
 * Ensure consistent float precision across tensors, NumPy arrays, and prior
   distributions.

meridian/model/eda/eda_engine.py

@@ -19,6 +19,7 @@
 from collections.abc import Collection, Sequence
 import dataclasses
 import functools
+import textwrap
 import typing
 from typing import Protocol
 import warnings
@@ -2492,3 +2493,66 @@ def check_population_corr_raw_media(
         explanation=eda_constants.POPULATION_CORRELATION_RAW_MEDIA_INFO,
         check_name='check_population_corr_raw_media',
     )
+
+  def check_data_param_ratio(
+      self,
+  ) -> eda_outcome.EDAOutcome[eda_outcome.DataParameterRatioArtifact]:
+    """Checks the ratio of data points to model parameters.
+
+    Returns:
+      An EDAOutcome object with findings and result values.
+    """
+    n_geos = self._model_context.n_geos
+    n_times = self._model_context.n_times
+    n_knots = self._model_context.knot_info.n_knots
+    n_controls = self._model_context.n_controls
+    n_treatments = (
+        self._model_context.n_media_channels
+        + self._model_context.n_rf_channels
+        + self._model_context.n_organic_media_channels
+        + self._model_context.n_organic_rf_channels
+        + self._model_context.n_non_media_channels
+    )
+
+    n_parameters = (n_geos - 1) + n_knots + n_controls + n_treatments
+    n_data_points = n_geos * n_times
+    ratio = n_data_points / n_parameters if n_parameters > 0 else float('inf')
+
+    artifact = eda_outcome.DataParameterRatioArtifact(
+        level=eda_outcome.AnalysisLevel.OVERALL,
+        n_parameters=n_parameters,
+        n_data_points=n_data_points,
+        ratio=ratio,
+    )
+
+    explanation = textwrap.dedent(f"""\
+    As a rough guidance, please review the ratio of data points to
+    parameters, where
+    * the number of data points = n_geos * n_times,
+    * the number of parameters = (n_geos-1) + n_knots + n_controls + n_treatments.\n
+    A very small ratio indicates insufficient data for estimation.
+    In that case, consider dropping or combining channels,
+    or reducing the number of knots with `knots` argument in `ModelSpec`.
+    For more details, please refer to this documentation page:
+    https://developers.google.com/meridian/docs/pre-modeling/amount-data-needed.\n
+    This ratio is {ratio:.2f} for your dataset, where
+    * n_geos = {n_geos}
+    * n_times = {n_times}
+    * n_knots = {n_knots}
+    * n_controls = {n_controls}
+    * n_treatments = {n_treatments}.""")
+
+    findings = [
+        eda_outcome.EDAFinding(
+            severity=eda_outcome.EDASeverity.INFO,
+            explanation=explanation,
+            finding_cause=eda_outcome.FindingCause.NONE,
+            associated_artifact=artifact,
+        )
+    ]
+
+    return eda_outcome.EDAOutcome(
+        check_type=eda_outcome.EDACheckType.DATA_ADEQUACY,
+        findings=findings,
+        analysis_artifacts=[artifact],
+    )

meridian/model/eda/eda_engine_test.py

@@ -21,6 +21,7 @@
 from meridian import constants
 from meridian.backend import test_utils
 from meridian.model import context
+from meridian.model import knots
 from meridian.model import model_test_data
 from meridian.model import spec as model_spec
 from meridian.model.eda import constants as eda_constants
@@ -8060,6 +8061,76 @@ def test_check_population_corr_raw_media_returns_info_finding(self):
     with self.subTest("finding_associated_artifact"):
       self.assertIs(finding.associated_artifact, outcome.analysis_artifacts[0])
 
+  def test_check_data_param_ratio(self):
+    mock_model_context = mock.create_autospec(
+        context.ModelContext, instance=True, spec_set=True
+    )
+
+    mock_model_context.n_geos = 10
+    mock_model_context.n_times = 20
+    mock_model_context.n_controls = 2
+    mock_model_context.n_media_channels = 3
+    mock_model_context.n_rf_channels = 1
+    mock_model_context.n_organic_media_channels = 0
+    mock_model_context.n_organic_rf_channels = 0
+    mock_model_context.n_non_media_channels = 0
+    n_knots = 5
+    mock_model_context.knot_info = knots.KnotInfo(
+        n_knots=n_knots, knot_locations=np.array([1]), weights=np.array([1])
+    )
+    mock.seal(mock_model_context)
+
+    engine = eda_engine.EDAEngine(model_context=mock_model_context)
+    outcome = engine.check_data_param_ratio()
+
+    expected_n_treatments = (
+        mock_model_context.n_media_channels
+        + mock_model_context.n_rf_channels
+        + mock_model_context.n_organic_media_channels
+        + mock_model_context.n_organic_rf_channels
+        + mock_model_context.n_non_media_channels
+    )
+    expected_n_parameters = (
+        (mock_model_context.n_geos - 1)
+        + n_knots
+        + mock_model_context.n_controls
+        + expected_n_treatments
+    )
+    expected_n_data_points = (
+        mock_model_context.n_geos * mock_model_context.n_times
+    )
+    expected_ratio = expected_n_data_points / expected_n_parameters
+
+    with self.subTest(name="check_type"):
+      self.assertEqual(
+          outcome.check_type, eda_outcome.EDACheckType.DATA_ADEQUACY
+      )
+
+    with self.subTest(name="finding_details"):
+      self.assertLen(outcome.findings, 1)
+      finding = outcome.findings[0]
+      self.assertEqual(finding.severity, eda_outcome.EDASeverity.INFO)
+      self.assertIn(
+          "As a rough guidance, please review the ratio of data points to",
+          finding.explanation,
+      )
+      self.assertIn(
+          f"This ratio is {expected_ratio:.2f} for your dataset",
+          finding.explanation,
+      )
+      self.assertIn(
+          f"* n_treatments = {expected_n_treatments}.",
+          finding.explanation,
+      )
+
+    with self.subTest(name="artifact_details"):
+      self.assertLen(outcome.analysis_artifacts, 1)
+      artifact = outcome.analysis_artifacts[0]
+      self.assertIsInstance(artifact, eda_outcome.DataParameterRatioArtifact)
+      self.assertEqual(artifact.n_parameters, expected_n_parameters)
+      self.assertEqual(artifact.n_data_points, expected_n_data_points)
+      self.assertAlmostEqual(artifact.ratio, expected_ratio)
+
 
 class HelpersTest(test_utils.MeridianTestCase):
 
@@ -8174,6 +8245,5 @@ def test_get_triangle_corr_mat_with_geo(self, lower, expected_values):
     actual = eda_engine.get_triangle_corr_mat(da, lower=lower)
     np.testing.assert_allclose(actual.values, expected_values)
 
-
 if __name__ == "__main__":
   absltest.main()

meridian/model/eda/eda_outcome.py

@@ -23,10 +23,10 @@
 
 __all__ = (
     "EDASeverity",
-    "EDAFinding",
+    "FindingCause",
     "AnalysisLevel",
     "AnalysisArtifact",
-    "FindingCause",
+    "EDAFinding",
     "PairwiseCorrArtifact",
     "StandardDeviationArtifact",
     "VIFArtifact",
@@ -35,8 +35,8 @@
     "VariableGeoTimeCollinearityArtifact",
     "PopulationCorrelationArtifact",
     "PriorProbabilityArtifact",
+    "DataParameterRatioArtifact",
     "EDACheckType",
-    "ArtifactType",
     "EDAOutcome",
     "CriticalCheckEDAOutcomes",
 )
@@ -260,6 +260,21 @@ class PriorProbabilityArtifact(AnalysisArtifact):
   mean_prior_contribution_da: xr.DataArray
 
 
+@dataclasses.dataclass(frozen=True)
+class DataParameterRatioArtifact(AnalysisArtifact):
+  """Artifact for model complexity check.
+
+  Attributes:
+    n_parameters: The number of parameters in the model.
+    n_data_points: The number of data points.
+    ratio: The ratio of data points to parameters.
+  """
+
+  n_parameters: int
+  n_data_points: int
+  ratio: float
+
+
 @enum.unique
 class EDACheckType(enum.Enum):
   """Enumeration for the type of an EDA check."""
@@ -272,6 +287,7 @@ class EDACheckType(enum.Enum):
   VARIABLE_GEO_TIME_COLLINEARITY = enum.auto()
   POPULATION_CORRELATION = enum.auto()
   PRIOR_PROBABILITY = enum.auto()
+  DATA_ADEQUACY = enum.auto()
 
 
 ArtifactType = typing.TypeVar("ArtifactType", bound=AnalysisArtifact)