From 95394e3b8da686015ffa7ed952de28c62f2ec1f6 Mon Sep 17 00:00:00 2001
From: Rieman Li <riemanli@google.com>
Date: Tue, 7 Jul 2026 18:42:17 -0700
Subject: [PATCH] Refactor `_generate_prior_calibration_chart_spec`

PiperOrigin-RevId: 944207623
---
 meridian/model/eda/meridian_eda.py      | 37 +++++++++++++++----------
 meridian/model/eda/meridian_eda_test.py |  1 +
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/meridian/model/eda/meridian_eda.py b/meridian/model/eda/meridian_eda.py
index c985f9289..e25984aa7 100644
--- a/meridian/model/eda/meridian_eda.py
+++ b/meridian/model/eda/meridian_eda.py
@@ -35,7 +35,6 @@
 import pandas as pd
 import xarray as xr
 
-
 if TYPE_CHECKING:
   from meridian.model import model  # pylint: disable=g-bad-import-order,g-import-not-at-top
 
@@ -1968,14 +1967,8 @@ def _generate_prior_specifications_card(
     [artifact] = self._dataset_level_prior_check_outcome.get_overall_artifacts()
     prior_probability = artifact.prior_negative_baseline_prob
 
-    prior_chart = formatter.ChartSpec(
-        id=eda_constants.PRIOR_CHART_ID,
-        chart_json=self.plot_prior_mean().to_json(),
-        infos=[
-            f'{eda_constants.PRIOR_PROBABILITY_REPORT_INFO}Prior Probability of'
-            f' negative baseline: {prior_probability}'
-        ],
-    )
+    prior_chart = self._generate_prior_mean_chart_spec(prior_probability)
+    chart_specs = [prior_chart]
 
     return (
         formatter.create_card_html(
@@ -1984,11 +1977,24 @@ def _generate_prior_specifications_card(
                 id=eda_constants.PRIOR_SPECIFICATIONS_CARD_ID,
                 title=eda_constants.PRIOR_SPECIFICATIONS_CARD_TITLE,
             ),
-            chart_specs=[prior_chart],
+            chart_specs=chart_specs,
         ),
         _initialize_severity_counts(),
     )
 
+  def _generate_prior_mean_chart_spec(
+      self, prior_probability: float
+  ) -> formatter.ChartSpec:
+    """Generates the ChartSpec for the default prior mean chart."""
+    return formatter.ChartSpec(
+        id=eda_constants.PRIOR_CHART_ID,
+        chart_json=self.plot_prior_mean().to_json(),
+        infos=[
+            f'{eda_constants.PRIOR_PROBABILITY_REPORT_INFO}Prior Probability of'
+            f' negative baseline: {prior_probability}'
+        ],
+    )
+
   def _validate_and_get_geos_to_plot(self, geos: Geos) -> Sequence[str]:
     """Validates and returns the geos to plot."""
     is_national = self._meridian.is_national
@@ -2274,10 +2280,13 @@ def _get_plot_data_for_heatmap(
   # Maintain the original order of variables.
   plot_variables = [v for v in all_variables if v in selected_vars]
 
-  return df[
-      df[eda_constants.VARIABLE_1].isin(plot_variables)
-      & df[eda_constants.VARIABLE_2].isin(plot_variables)
-  ], plot_variables
+  return (
+      df[
+          df[eda_constants.VARIABLE_1].isin(plot_variables)
+          & df[eda_constants.VARIABLE_2].isin(plot_variables)
+      ],
+      plot_variables,
+  )
 
 
 def _format_explanation_for_html(explanation: str) -> str:
diff --git a/meridian/model/eda/meridian_eda_test.py b/meridian/model/eda/meridian_eda_test.py
index a0f5bd131..b526dc5a9 100644
--- a/meridian/model/eda/meridian_eda_test.py
+++ b/meridian/model/eda/meridian_eda_test.py
@@ -34,6 +34,7 @@
 from meridian.model.eda import eda_spec
 from meridian.model.eda import meridian_eda
 from meridian.model.eda import sampling_eda_engine
+
 import numpy as np
 import pandas as pd
 import xarray as xr