Encapsulate GoF metrics into a dataclass

PiperOrigin-RevId: 864447565

Author: lukmaz

meridian/analysis/review/checks.py

@@ -186,7 +186,7 @@ def _set_metrics_from_gof_dataframe(
     metrics: MutableMapping[str, float],
     gof_df: pd.DataFrame,
     geo_granularity: str,
-    suffix: str | None = None,
+    suffix: str,
 ) -> None:
   """Sets the `metrics` variable of the GoodnessOfFitCheckResult.
 
@@ -200,31 +200,23 @@ def _set_metrics_from_gof_dataframe(
       holdout set is not used) of filtered to a single evaluation set ("all",
       "train", or "test").
     geo_granularity: The geo granularity of the data ("geo" or "national").
-    suffix: A suffix to add to the metric names (e.g., "all", "train", "test").
-      If None, the metrics are added without a suffix.
+    suffix: A suffix to add to the metric names (e.g., "_train", "_test").
   """
   gof_metrics_pivoted = gof_df.pivot(
       index=constants.GEO_GRANULARITY,
       columns=constants.METRIC,
       values=constants.VALUE,
   )
   gof_metrics_series = gof_metrics_pivoted.loc[geo_granularity]
-  if suffix is not None:
-    metrics[f"{review_constants.R_SQUARED}_{suffix}"] = gof_metrics_series[
-        constants.R_SQUARED
-    ]
-    metrics[f"{review_constants.MAPE}_{suffix}"] = gof_metrics_series[
-        constants.MAPE
-    ]
-    metrics[f"{review_constants.WMAPE}_{suffix}"] = gof_metrics_series[
-        constants.WMAPE
-    ]
-  else:
-    metrics[review_constants.R_SQUARED] = gof_metrics_series[
-        constants.R_SQUARED
-    ]
-    metrics[review_constants.MAPE] = gof_metrics_series[constants.MAPE]
-    metrics[review_constants.WMAPE] = gof_metrics_series[constants.WMAPE]
+  metrics[f"{review_constants.R_SQUARED}{suffix}"] = gof_metrics_series[
+      constants.R_SQUARED
+  ]
+  metrics[f"{review_constants.MAPE}{suffix}"] = gof_metrics_series[
+      constants.MAPE
+  ]
+  metrics[f"{review_constants.WMAPE}{suffix}"] = gof_metrics_series[
+      constants.WMAPE
+  ]
 
 
 class GoodnessOfFitCheck(
@@ -243,7 +235,7 @@ def run(self) -> results.GoodnessOfFitCheckResult:
     gof_metrics = gof_df[gof_df[constants.GEO_GRANULARITY] == geo_granularity]
     is_holdout = constants.EVALUATION_SET_VAR in gof_df.columns
 
-    metrics = {}
+    metrics_dict = {}
     case = results.GoodnessOfFitCases.PASS
 
     if is_holdout:
@@ -256,28 +248,70 @@ def run(self) -> results.GoodnessOfFitCheckResult:
             gof_metrics[constants.EVALUATION_SET_VAR] == evaluation_set
         ]
         _set_metrics_from_gof_dataframe(
-            metrics=metrics,
+            metrics=metrics_dict,
             gof_df=set_metrics,
             geo_granularity=geo_granularity,
             suffix=suffix,
         )
-        if metrics[f"{review_constants.R_SQUARED}_{suffix}"] <= 0:
+        if metrics_dict[f"{review_constants.R_SQUARED}{suffix}"] <= 0:
           case = results.GoodnessOfFitCases.REVIEW
+      return results.GoodnessOfFitCheckResult(
+          case=case,
+          metrics=results.GoodnessOfFitMetrics(
+              r_squared=metrics_dict[
+                  f"{review_constants.R_SQUARED}{review_constants.ALL_SUFFIX}"
+              ],
+              mape=metrics_dict[
+                  f"{review_constants.MAPE}{review_constants.ALL_SUFFIX}"
+              ],
+              wmape=metrics_dict[
+                  f"{review_constants.WMAPE}{review_constants.ALL_SUFFIX}"
+              ],
+              r_squared_train=metrics_dict[
+                  f"{review_constants.R_SQUARED}{review_constants.TRAIN_SUFFIX}"
+              ],
+              mape_train=metrics_dict[
+                  f"{review_constants.MAPE}{review_constants.TRAIN_SUFFIX}"
+              ],
+              wmape_train=metrics_dict[
+                  f"{review_constants.WMAPE}{review_constants.TRAIN_SUFFIX}"
+              ],
+              r_squared_test=metrics_dict[
+                  f"{review_constants.R_SQUARED}{review_constants.TEST_SUFFIX}"
+              ],
+              mape_test=metrics_dict[
+                  f"{review_constants.MAPE}{review_constants.TEST_SUFFIX}"
+              ],
+              wmape_test=metrics_dict[
+                  f"{review_constants.WMAPE}{review_constants.TEST_SUFFIX}"
+              ],
+          ),
+          is_holdout=is_holdout,
+      )
     else:
       _set_metrics_from_gof_dataframe(
-          metrics=metrics,
+          metrics=metrics_dict,
           gof_df=gof_metrics,
           geo_granularity=geo_granularity,
-          suffix=None,
+          suffix=review_constants.ALL_SUFFIX,
       )
-      if metrics[review_constants.R_SQUARED] <= 0:
+      if metrics_dict[review_constants.R_SQUARED] <= 0:
         case = results.GoodnessOfFitCases.REVIEW
-
-    return results.GoodnessOfFitCheckResult(
-        case=case,
-        metrics=metrics,
-        is_holdout=is_holdout,
-    )
+      return results.GoodnessOfFitCheckResult(
+          case=case,
+          metrics=results.GoodnessOfFitMetrics(
+              r_squared=metrics_dict[
+                  f"{review_constants.R_SQUARED}{review_constants.ALL_SUFFIX}"
+              ],
+              mape=metrics_dict[
+                  f"{review_constants.MAPE}{review_constants.ALL_SUFFIX}"
+              ],
+              wmape=metrics_dict[
+                  f"{review_constants.WMAPE}{review_constants.ALL_SUFFIX}"
+              ],
+          ),
+          is_holdout=is_holdout,
+      )
 
 
 # ==============================================================================

meridian/analysis/review/constants.py

@@ -32,9 +32,9 @@
 R_SQUARED = "r_squared"
 MAPE = "mape"
 WMAPE = "wmape"
-ALL_SUFFIX = "all"
-TRAIN_SUFFIX = "train"
-TEST_SUFFIX = "test"
+ALL_SUFFIX = ""
+TRAIN_SUFFIX = "_train"
+TEST_SUFFIX = "_test"
 EVALUATION_SET_SUFFIXES = (ALL_SUFFIX, TRAIN_SUFFIX, TEST_SUFFIX)
 MEAN = "mean"
 VARIANCE = "variance"

meridian/analysis/review/results.py

@@ -345,157 +345,77 @@ def __init__(
     super().__init__(status, message_template, recommendation)
 
 
+@dataclasses.dataclass(frozen=True)
+class GoodnessOfFitMetrics:
+  """The metrics for the Goodness of Fit Check."""
+
+  r_squared: float
+  mape: float
+  wmape: float
+  r_squared_train: float | None = None
+  mape_train: float | None = None
+  wmape_train: float | None = None
+  r_squared_test: float | None = None
+  mape_test: float | None = None
+  wmape_test: float | None = None
+
+
 @dataclasses.dataclass(frozen=True)
 class GoodnessOfFitCheckResult(CheckResult):
   """The immutable result of the Goodness of Fit Check."""
 
   case: GoodnessOfFitCases
-  metrics: Mapping[str, float]
+  metrics: GoodnessOfFitMetrics
   is_holdout: bool = False
 
   def __post_init__(self):
     if self.is_holdout:
-      required_keys = []
-      for suffix in [
-          constants.ALL_SUFFIX,
-          constants.TRAIN_SUFFIX,
-          constants.TEST_SUFFIX,
-      ]:
-        required_keys.extend([
-            f"{constants.R_SQUARED}_{suffix}",
-            f"{constants.MAPE}_{suffix}",
-            f"{constants.WMAPE}_{suffix}",
-        ])
-      if any(key not in self.metrics for key in required_keys):
+      if any(
+          metric is None
+          for metric in (
+              self.metrics.r_squared_train,
+              self.metrics.mape_train,
+              self.metrics.wmape_train,
+              self.metrics.r_squared_test,
+              self.metrics.mape_test,
+              self.metrics.wmape_test,
+          )
+      ):
         raise ValueError(
             "The message template is missing required formatting arguments for"
-            f" holdout case. Required keys: {required_keys}. Metrics:"
+            " holdout case. Required keys: r_squared_train, mape_train,"
+            " wmape_train, r_squared_test, mape_test, wmape_test. Metrics:"
             f" {self.metrics}."
         )
-    elif any(
-        key not in self.metrics
-        for key in (
-            constants.R_SQUARED,
-            constants.MAPE,
-            constants.WMAPE,
-        )
-    ):
-      raise ValueError(
-          "The message template is missing required formatting arguments:"
-          " r_squared, mape, wmape. Metrics:"
-          f" {self.metrics}."
-      )
-
-  @property
-  def r_squared(self) -> float | None:
-    """The R-squared metric."""
-    return self.metrics[constants.R_SQUARED] if not self.is_holdout else None
-
-  @property
-  def mape(self) -> float | None:
-    """The MAPE metric."""
-    return self.metrics[constants.MAPE] if not self.is_holdout else None
-
-  @property
-  def wmape(self) -> float | None:
-    """The wMAPE metric."""
-    return self.metrics[constants.WMAPE] if not self.is_holdout else None
-
-  @property
-  def r_squared_all(self) -> float | None:
-    """The R-squared metric for all data."""
-    return (
-        self.metrics[f"{constants.R_SQUARED}_{constants.ALL_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def mape_all(self) -> float | None:
-    """The MAPE metric for all data."""
-    return (
-        self.metrics[f"{constants.MAPE}_{constants.ALL_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def wmape_all(self) -> float | None:
-    """The wMAPE metric for all data."""
-    return (
-        self.metrics[f"{constants.WMAPE}_{constants.ALL_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def r_squared_train(self) -> float | None:
-    """The R-squared metric for train data."""
-    return (
-        self.metrics[f"{constants.R_SQUARED}_{constants.TRAIN_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def mape_train(self) -> float | None:
-    """The MAPE metric for train data."""
-    return (
-        self.metrics[f"{constants.MAPE}_{constants.TRAIN_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def wmape_train(self) -> float | None:
-    """The wMAPE metric for train data."""
-    return (
-        self.metrics[f"{constants.WMAPE}_{constants.TRAIN_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def r_squared_test(self) -> float | None:
-    """The R-squared metric for test data."""
-    return (
-        self.metrics[f"{constants.R_SQUARED}_{constants.TEST_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def mape_test(self) -> float | None:
-    """The MAPE metric for test data."""
-    return (
-        self.metrics[f"{constants.MAPE}_{constants.TEST_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
-
-  @property
-  def wmape_test(self) -> float | None:
-    """The wMAPE metric for test data."""
-    return (
-        self.metrics[f"{constants.WMAPE}_{constants.TEST_SUFFIX}"]
-        if self.is_holdout
-        else None
-    )
 
   @property
   def details(self) -> Mapping[str, Any]:
     """The check result details."""
-    return self.metrics
+    return {
+        f"{constants.R_SQUARED}{constants.ALL_SUFFIX}": self.metrics.r_squared,
+        f"{constants.MAPE}{constants.ALL_SUFFIX}": self.metrics.mape,
+        f"{constants.WMAPE}{constants.ALL_SUFFIX}": self.metrics.wmape,
+        f"{constants.R_SQUARED}{constants.TRAIN_SUFFIX}": (
+            self.metrics.r_squared_train
+        ),
+        f"{constants.MAPE}{constants.TRAIN_SUFFIX}": self.metrics.mape_train,
+        f"{constants.WMAPE}{constants.TRAIN_SUFFIX}": self.metrics.wmape_train,
+        f"{constants.R_SQUARED}{constants.TEST_SUFFIX}": (
+            self.metrics.r_squared_test
+        ),
+        f"{constants.MAPE}{constants.TEST_SUFFIX}": self.metrics.mape_test,
+        f"{constants.WMAPE}{constants.TEST_SUFFIX}": self.metrics.wmape_test,
+    }
 
   @property
   def recommendation(self) -> str:
     """The check result message."""
     if self.is_holdout:
       report_str = (
-          "R-squared = {r_squared_all:.4f} (All),"
+          "R-squared = {r_squared:.4f} (All),"
           " {r_squared_train:.4f} (Train), {r_squared_test:.4f} (Test); MAPE"
-          " = {mape_all:.4f} (All), {mape_train:.4f} (Train),"
-          " {mape_test:.4f} (Test); wMAPE = {wmape_all:.4f} (All),"
+          " = {mape:.4f} (All), {mape_train:.4f} (Train),"
+          " {mape_test:.4f} (Test); wMAPE = {wmape:.4f} (All),"
           " {wmape_train:.4f} (Train), {wmape_test:.4f} (Test)".format(
               **self.details
           )