Refactor model fit plot to include interactive tooltips and hover effects

PiperOrigin-RevId: 876095304

Author: The Meridian Authors

meridian/analysis/visualizer.py

@@ -481,53 +481,88 @@ def plot_model_fit(
       y_axis_label = summary_text.KPI_LABEL
     else:
       y_axis_label = summary_text.REVENUE_LABEL
-    plot = (
-        alt.Chart(model_fit_df, width=c.VEGALITE_FACET_EXTRA_LARGE_WIDTH)
-        .mark_line()
-        .encode(
-            x=alt.X(
-                f'{c.TIME}:T',
-                title='Time period',
-                axis=alt.Axis(
-                    format=c.QUARTER_FORMAT,
-                    grid=False,
-                    tickCount=8,
-                    domainColor=c.GREY_300,
-                ),
-            ),
-            y=alt.Y(
-                f'{c.MEAN}:Q',
-                title=y_axis_label,
-                axis=alt.Axis(
-                    ticks=False,
-                    domain=False,
-                    tickCount=5,
-                    labelPadding=c.PADDING_10,
-                    labelExpr=formatter.compact_number_expr(),
-                    **formatter.Y_AXIS_TITLE_CONFIG,
-                ),
+
+    base = alt.Chart(model_fit_df, width=c.VEGALITE_FACET_EXTRA_LARGE_WIDTH)
+
+    hover = alt.selection_point(
+        fields=[c.TIME],
+        nearest=True,
+        on='pointerover',
+        empty=False,
+        clear='pointerout',
+    )
+
+    lines = base.mark_line().encode(
+        x=alt.X(
+            f'{c.TIME}:T',
+            title='Time period',
+            axis=alt.Axis(
+                format=c.QUARTER_FORMAT,
+                grid=False,
+                tickCount=8,
+                domainColor=c.GREY_300,
             ),
-            color=alt.Color(
-                'type:N', scale=alt.Scale(domain=domain, range=colors)
+        ),
+        y=alt.Y(
+            f'{c.MEAN}:Q',
+            title=y_axis_label,
+            axis=alt.Axis(
+                ticks=False,
+                domain=False,
+                tickCount=5,
+                labelPadding=c.PADDING_10,
+                labelExpr=formatter.compact_number_expr(),
+                **formatter.Y_AXIS_TITLE_CONFIG,
             ),
+        ),
+        color=alt.Color('type:N', scale=alt.Scale(domain=domain, range=colors)),
+    )
+
+    tooltip = [alt.Tooltip(f'{c.TIME}:T', title='Time period')]
+    for field in domain:
+      tooltip.append(
+          alt.Tooltip(
+              f'{field}:Q',
+              title=field.title(),
+              format=',.0f',
+          )
+      )
+
+    tooltips = (
+        base.transform_pivot(c.TYPE, value=c.MEAN, groupby=[c.TIME])
+        .mark_rule(opacity=0)
+        .encode(
+            x=f'{c.TIME}:T',
+            opacity=alt.condition(hover, alt.value(0.3), alt.value(0)),
+            tooltip=tooltip,
         )
+        .add_params(hover)
     )
 
+    points = base.mark_circle(size=c.MARK_CIRCLE_SIZE, filled=False).encode(
+        x=f'{c.TIME}:T',
+        y=f'{c.MEAN}:Q',
+        color=alt.Color(
+            'type:N',
+            scale=alt.Scale(domain=domain, range=colors),
+            legend=None,
+        ),
+        opacity=alt.condition(hover, alt.value(1), alt.value(0)),
+    )
+
+    plot = alt.layer(lines, tooltips, points)
+
     if include_ci:
       # Only add a confidence interval area for the modeled data.
-      confidence_band = (
-          alt.Chart(model_fit_df)
-          .mark_area(opacity=0.3)
-          .encode(
-              x=f'{c.TIME}:T',
-              y=f'{c.CI_HI}:Q',
-              y2=f'{c.CI_LO}:Q',
-              color=alt.Color(
-                  'type:N',
-                  scale=alt.Scale(domain=[domain[0]], range=[colors[0]]),
-                  legend=None,
-              ),
-          )
+      confidence_band = base.mark_area(opacity=0.3).encode(
+          x=f'{c.TIME}:T',
+          y=f'{c.CI_HI}:Q',
+          y2=f'{c.CI_LO}:Q',
+          color=alt.Color(
+              'type:N',
+              scale=alt.Scale(domain=[domain[0]], range=[colors[0]]),
+              legend=None,
+          ),
       )
       plot = (plot + confidence_band).resolve_scale(color=c.INDEPENDENT)
 

meridian/analysis/visualizer_test.py

@@ -703,16 +703,16 @@ def test_model_fit_plots_expected_ci(self):
 
     self.assertIsInstance(plot, alt.LayerChart)
     self.assertEqual(
-        plot.layer[1].encoding.color["scale"]["domain"], [c.EXPECTED]
+        plot.layer[3].encoding.color["scale"]["domain"], [c.EXPECTED]
     )
-    self.assertEqual(plot.layer[1].encoding.y.shorthand, f"{c.CI_HI}:Q")
-    self.assertEqual(plot.layer[1].encoding.y2.shorthand, f"{c.CI_LO}:Q")
+    self.assertEqual(plot.layer[3].encoding.y.shorthand, f"{c.CI_HI}:Q")
+    self.assertEqual(plot.layer[3].encoding.y2.shorthand, f"{c.CI_LO}:Q")
 
   def test_model_fit_plots_no_ci(self):
     plot = self.model_fit_kpi_type_revenue.plot_model_fit(include_ci=False)
-    self.assertIsInstance(plot, alt.Chart)
-    self.assertEqual(plot.encoding.x.shorthand, f"{c.TIME}:T")
-    self.assertEqual(plot.encoding.y.shorthand, f"{c.MEAN}:Q")
+    self.assertIsInstance(plot, alt.LayerChart)
+    self.assertEqual(plot.layer[0].encoding.x.shorthand, f"{c.TIME}:T")
+    self.assertEqual(plot.layer[0].encoding.y.shorthand, f"{c.MEAN}:Q")
 
   def test_model_fit_axis_encoding(self):
     plot = self.model_fit_kpi_type_revenue.plot_model_fit()
@@ -737,6 +737,43 @@ def test_model_fit_axis_encoding(self):
         | formatter.Y_AXIS_TITLE_CONFIG,
     )
 
+  def test_model_fit_tooltip_encoding(self):
+    plot = self.model_fit_kpi_type_revenue.plot_model_fit()
+    tooltip_layer = plot.layer[1]
+    self.assertEqual(tooltip_layer.mark.type, "rule")
+    self.assertEqual(tooltip_layer.encoding.x.shorthand, f"{c.TIME}:T")
+    self.assertFalse(tooltip_layer.encoding.opacity["condition"].empty)
+    expected_tooltip = [
+        alt.Tooltip(f"{c.TIME}:T", title="Time period"),
+        alt.Tooltip(f"{c.EXPECTED}:Q", title="Expected", format=",.0f"),
+        alt.Tooltip(f"{c.ACTUAL}:Q", title="Actual", format=",.0f"),
+        alt.Tooltip(f"{c.BASELINE}:Q", title="Baseline", format=",.0f"),
+    ]
+    self.assertEqual(tooltip_layer.encoding.tooltip, expected_tooltip)
+
+  def test_model_fit_points_encoding(self):
+    plot = self.model_fit_kpi_type_revenue.plot_model_fit()
+    points_layer = plot.layer[2]
+
+    with self.subTest("test_mark_properties"):
+      self.assertEqual(points_layer.mark.type, "circle")
+      self.assertFalse(points_layer.mark.filled)
+      self.assertEqual(points_layer.mark.size, c.MARK_CIRCLE_SIZE)
+
+    with self.subTest("test_encoding_shorthands"):
+      self.assertEqual(points_layer.encoding.x.shorthand, f"{c.TIME}:T")
+      self.assertEqual(points_layer.encoding.y.shorthand, f"{c.MEAN}:Q")
+
+    with self.subTest("test_color_legend"):
+      self.assertIsNone(points_layer.encoding.color["legend"])
+
+  def test_model_fit_selection_points(self):
+    plot = self.model_fit_kpi_type_revenue.plot_model_fit()
+    selection_points = plot.params[0].select
+    self.assertTrue(selection_points["nearest"])
+    self.assertEqual(selection_points["on"], "pointerover")
+    self.assertEqual(selection_points["clear"], "pointerout")
+
   def test_model_fit_correct_config(self):
     plot = self.model_fit_kpi_type_revenue.plot_model_fit()
     self.assertEqual(plot.config.axis.to_dict(), formatter.TEXT_CONFIG)

meridian/constants.py

@@ -758,6 +758,7 @@
 INDEPENDENT = 'independent'
 RESPONSE_CURVE_STEP_SIZE = 0.01
 OUTLIER_CLIP_FACTOR = 1.2
+MARK_CIRCLE_SIZE = 36
 
 
 # Font names.