feat(viz): per-facet shading for faceted regression-fit plots

CHANGELOG.md

@@ -4,6 +4,13 @@
 
 Full parity with the R `moderndive` and `infer` packages.
 
+- **Per-facet shading for regression-fit plots**: `shade_p_value` and
+  `shade_confidence_interval` now accept per-term values (an observed `FitResult`,
+  a `term`-keyed CI/p-value table, or a dict) so each facet of a faceted
+  `visualize_fit()` plot is shaded from its own observed statistic / interval —
+  in both the plotly and plotnine engines. (Previously shading was scalar-only
+  and couldn't shade per facet, so faceted multiple-regression inference plots
+  rendered without the overlay.)
 - **Dual-engine plotting**: every plotting function (`visualize`,
   `shade_p_value`, `shade_confidence_interval`, `pairplot`, and the new model
   plots) now takes `engine="plotly"` (the new default) or `engine="plotnine"`.

moderndive/infer/core.py

@@ -422,10 +422,12 @@ def get_p_value(self, obs_stat: FitResult, direction: str = "two-sided") -> pl.D
 
         return get_fit_p_value(self, obs_stat=obs_stat, direction=direction)
 
-    def visualize(self, bins: int = 20, *, engine: str = "plotly"):
+    def visualize(self, bins: int = 20, *, engine: str = "plotly", shade_pvalue=None, shade_ci=None):
         from .viz import visualize_fit
 
-        return visualize_fit(self, bins=bins, engine=engine)
+        return visualize_fit(
+            self, bins=bins, engine=engine, shade_pvalue=shade_pvalue, shade_ci=shade_ci
+        )
 
 
 def specify(

moderndive/infer/viz/__init__.py

@@ -40,14 +40,22 @@
 
 @dataclass(frozen=True)
 class ShadeSpec:
-    """Engine-neutral description of a shading layer (p-value tail or CI)."""
+    """Engine-neutral description of a shading layer (p-value tail or CI).
+
+    For a single-panel plot the scalar fields are used. For a faceted
+    regression-fit plot (:func:`visualize_fit`), ``per_term`` holds one value per
+    term so each facet is shaded independently — a tuple of ``(term, payload)``
+    pairs where ``payload`` is the observed statistic (p-value) or a
+    ``(lower, upper)`` pair (confidence interval).
+    """
 
     kind: str  # "p_value" | "confidence_interval"
     obs_stat: float | None = None
     direction: str | None = None
     lower: float | None = None
     upper: float | None = None
     color: str | None = None
+    per_term: tuple | None = None
 
 
 class InferPlot:
@@ -58,11 +66,15 @@ class InferPlot:
     raw ``ggplot`` is also available via :attr:`gg`.
     """
 
-    def __init__(self, figure, engine: str):
+    def __init__(self, figure, engine: str, terms: list[str] | None = None):
         self.figure = figure
         self.engine = engine
+        # Facet terms for a regression-fit plot (enables per-facet shading).
+        self.terms = terms
 
     def __add__(self, other):
+        if isinstance(other, ShadeSpec) and other.per_term is not None:
+            return self._add_per_facet(other)
         if self.engine == "plotnine":
             from . import _plotnine as P
 
@@ -72,19 +84,34 @@ def __add__(self, other):
                     if other.kind == "p_value"
                     else P.shade_ci_layers(other)
                 )
-                return InferPlot(self.figure + layers, "plotnine")
+                return InferPlot(self.figure + layers, "plotnine", self.terms)
             # Any other plotnine object/layer/list.
-            return InferPlot(self.figure + other, "plotnine")
+            return InferPlot(self.figure + other, "plotnine", self.terms)
 
         from . import _plotly as PX
 
         if isinstance(other, ShadeSpec):
-            return InferPlot(PX.apply_shade_px(self.figure, other), "plotly")
+            return InferPlot(PX.apply_shade_px(self.figure, other), "plotly", self.terms)
         raise TypeError(
             "Only a ShadeSpec (from shade_p_value/shade_confidence_interval) can be "
             "added to a plotly InferPlot."
         )
 
+    def _add_per_facet(self, spec: ShadeSpec):
+        """Apply per-term shading to a faceted regression-fit plot (both engines)."""
+        if self.terms is None:
+            raise TypeError(
+                "Per-term shading requires a faceted fit plot from visualize_fit(); "
+                "pass a per-term obs_stat/endpoints (a FitResult or a term-keyed table)."
+            )
+        if self.engine == "plotnine":
+            from . import _plotnine as P
+
+            return InferPlot(P.apply_fit_shade_gg(self.figure, spec, self.terms), "plotnine", self.terms)
+        from . import _plotly as PX
+
+        return InferPlot(PX.apply_fit_shade_px(self.figure, spec, self.terms), "plotly", self.terms)
+
     @property
     def gg(self):
         """The underlying plotnine ``ggplot`` (only when ``engine='plotnine'``)."""
@@ -166,16 +193,32 @@ def visualize(
     return plot
 
 
-def visualize_fit(fit, bins: int = 20, *, engine: str = "plotly") -> InferPlot:
-    """Faceted histogram of a regression fit distribution, one panel per term."""
+def visualize_fit(
+    fit, bins: int = 20, *, engine: str = "plotly", shade_pvalue=None, shade_ci=None
+) -> InferPlot:
+    """Faceted histogram of a regression fit distribution, one panel per term.
+
+    Pass ``shade_pvalue=``/``shade_ci=`` to shade each facet from per-term values
+    (a ``FitResult`` of observed estimates, or a term-keyed CI/p-value table), or
+    compose the same per-term :class:`ShadeSpec` with ``+``.
+    """
     engine = C.resolve_engine(engine)
+    terms = fit.data["term"].unique(maintain_order=True).to_list()
     if engine == "plotnine":
         from . import _plotnine as P
 
-        return InferPlot(P.visualize_fit_gg(fit, bins), engine)
-    from . import _plotly as PX
+        fig = P.visualize_fit_gg(fit, bins)
+    else:
+        from . import _plotly as PX
+
+        fig = PX.visualize_fit_px(fit, bins)
 
-    return InferPlot(PX.visualize_fit_px(fit, bins), engine)
+    plot = InferPlot(fig, engine, terms)
+    if shade_pvalue is not None:
+        plot = plot + _coerce_pvalue_spec(shade_pvalue)
+    if shade_ci is not None:
+        plot = plot + _coerce_ci_spec(shade_ci)
+    return plot
 
 
 def visualize_theoretical(theoretical, bins: int = 100, *, engine: str = "plotly") -> InferPlot:
@@ -195,20 +238,63 @@ def visualize_theoretical(theoretical, bins: int = 100, *, engine: str = "plotly
     return InferPlot(PX.density_curve_px(x, density, title), engine)
 
 
+def _per_term_obs(obs_stat) -> dict | None:
+    """Extract a ``{term: observed}`` mapping for per-facet p-value shading.
+
+    Accepts an observed ``FitResult`` (term/estimate), a term-keyed polars frame
+    (``term`` + ``estimate`` or ``stat``), or a dict. Returns ``None`` for a scalar.
+    """
+    import polars as pl
+
+    if isinstance(obs_stat, pl.DataFrame) and "term" in obs_stat.columns:
+        valcol = "estimate" if "estimate" in obs_stat.columns else "stat"
+        return {r["term"]: float(r[valcol]) for r in obs_stat.iter_rows(named=True)}
+    data = getattr(obs_stat, "data", None)
+    if isinstance(data, pl.DataFrame) and {"term", "estimate"} <= set(data.columns):
+        return {r["term"]: float(r["estimate"]) for r in data.iter_rows(named=True)}
+    if isinstance(obs_stat, dict):
+        return {k: float(v) for k, v in obs_stat.items()}
+    return None
+
+
+def _per_term_ci(endpoints) -> dict | None:
+    """Extract a ``{term: (lower, upper)}`` mapping for per-facet CI shading."""
+    import polars as pl
+
+    if isinstance(endpoints, pl.DataFrame) and "term" in endpoints.columns:
+        return {
+            r["term"]: (float(r["lower_ci"]), float(r["upper_ci"]))
+            for r in endpoints.iter_rows(named=True)
+        }
+    return None
+
+
 def shade_p_value(obs_stat, direction: str, *, color: str | None = None) -> ShadeSpec:
     """A p-value shading spec; add it to a ``visualize()`` plot with ``+``.
 
-    ``direction`` ∈ {right/greater, left/less, two-sided}.
+    ``direction`` ∈ {right/greater, left/less, two-sided}. For a faceted
+    :func:`visualize_fit` plot, pass a per-term ``obs_stat`` — an observed
+    ``FitResult``, a ``term``-keyed frame, or a dict — to shade each facet.
     """
-    return ShadeSpec(
-        kind="p_value", obs_stat=float(obs_stat), direction=direction, color=color
-    )
+    per = _per_term_obs(obs_stat)
+    if per is not None:
+        return ShadeSpec(
+            kind="p_value", direction=direction, color=color, per_term=tuple(sorted(per.items()))
+        )
+    return ShadeSpec(kind="p_value", obs_stat=float(obs_stat), direction=direction, color=color)
 
 
 def shade_confidence_interval(endpoints, color: str | None = None) -> ShadeSpec:
     """A confidence-interval shading spec; add it to a ``visualize()`` plot with ``+``.
 
-    ``endpoints`` is a CI DataFrame (``lower_ci``/``upper_ci``) or a ``(lower, upper)`` tuple.
+    ``endpoints`` is a CI DataFrame (``lower_ci``/``upper_ci``) or a ``(lower, upper)``
+    tuple. For a faceted :func:`visualize_fit` plot, pass a per-term CI table (with a
+    ``term`` column) to shade each facet from its own interval.
     """
+    per = _per_term_ci(endpoints)
+    if per is not None:
+        return ShadeSpec(
+            kind="confidence_interval", color=color, per_term=tuple(sorted(per.items()))
+        )
     lower, upper = C.ci_endpoints(endpoints)
     return ShadeSpec(kind="confidence_interval", lower=lower, upper=upper, color=color)

moderndive/infer/viz/_plotly.py

@@ -131,3 +131,64 @@ def apply_shade_px(fig, spec):
         out.add_vline(x=spec.lower, line={"color": color, "width": 2})
         out.add_vline(x=spec.upper, line={"color": color, "width": 2})
     return out
+
+
+def _subplot_xrange(fig, col: int) -> tuple[float, float]:
+    """Finite x-range of the histogram in subplot ``col`` (for clipping infinite shades)."""
+    axis = "x" if col == 1 else f"x{col}"
+    xs = []
+    for trace in fig.data:
+        tx = getattr(trace, "xaxis", None) or "x"
+        if tx == axis and getattr(trace, "x", None) is not None and len(trace.x):
+            xs.append(np.asarray(trace.x, dtype=float))
+    if not xs:
+        return -1.0, 1.0
+    allx = np.concatenate(xs)
+    lo, hi = float(np.min(allx)), float(np.max(allx))
+    pad = (hi - lo) * 0.05 or 1.0
+    return lo - pad, hi + pad
+
+
+def apply_fit_shade_px(fig, spec, terms):
+    """Per-facet shading for a faceted fit figure: shade each term's subplot via row/col."""
+    go = _go()
+    out = go.Figure(fig)
+    term_to_col = {t: i + 1 for i, t in enumerate(terms)}
+    per = dict(spec.per_term)
+
+    if spec.kind == "p_value":
+        for term, obs in per.items():
+            col = term_to_col.get(term)
+            if col is None:
+                continue
+            lo, hi = _subplot_xrange(out, col)
+            vlines, rects = C.pvalue_regions(obs, spec.direction)
+            for x, dashed in vlines:
+                out.add_vline(
+                    x=x,
+                    line={"color": C._OBS_COLOR, "width": 2, "dash": "dash" if dashed else "solid"},
+                    row=1,
+                    col=col,
+                )
+            for xmin, xmax in rects:
+                out.add_vrect(
+                    x0=_clip(xmin, lo, hi),
+                    x1=_clip(xmax, lo, hi),
+                    fillcolor=C._OBS_COLOR,
+                    opacity=0.3,
+                    line_width=0,
+                    row=1,
+                    col=col,
+                )
+    else:  # confidence_interval
+        color = spec.color or C._SHADE_COLOR
+        for term, (lower, upper) in per.items():
+            col = term_to_col.get(term)
+            if col is None:
+                continue
+            out.add_vrect(
+                x0=lower, x1=upper, fillcolor=color, opacity=0.3, line_width=0, row=1, col=col
+            )
+            out.add_vline(x=lower, line={"color": color, "width": 2}, row=1, col=col)
+            out.add_vline(x=upper, line={"color": color, "width": 2}, row=1, col=col)
+    return out

moderndive/infer/viz/_plotnine.py

@@ -10,6 +10,7 @@
     facet_wrap,
     geom_histogram,
     geom_line,
+    geom_rect,
     geom_vline,
     ggplot,
     labs,
@@ -102,3 +103,69 @@ def shade_ci_layers(spec) -> list:
         geom_vline(xintercept=spec.lower, color=color, size=1.0),
         geom_vline(xintercept=spec.upper, color=color, size=1.0),
     ]
+
+
+def _facet_rect(rows, fill: str):
+    """A geom_rect keyed by the ``term`` facet column, spanning each panel's height."""
+    return geom_rect(
+        aes(xmin="xmin", xmax="xmax", group="term"),
+        data=pd.DataFrame(rows),
+        ymin=-C._INF,
+        ymax=C._INF,
+        alpha=0.3,
+        fill=fill,
+        inherit_aes=False,
+    )
+
+
+def _facet_vlines(rows, color: str, dashed: bool):
+    extra = {"linetype": "dashed"} if dashed else {}
+    return geom_vline(
+        aes(xintercept="x", group="term"),
+        data=pd.DataFrame(rows),
+        color=color,
+        size=1.0,
+        inherit_aes=False,
+        **extra,
+    )
+
+
+def apply_fit_shade_gg(gg, spec, terms):
+    """Per-facet shading for a faceted fit plot: each ``term`` panel shaded on its own.
+
+    Shading data carries a ``term`` column matching ``facet_wrap("term")`` so each
+    layer lands only in its panel. ``inherit_aes=False`` keeps the histogram's
+    ``x="estimate"`` mapping from leaking into these layers.
+    """
+    per = dict(spec.per_term)
+    layers = []
+    if spec.kind == "p_value":
+        solid, dashed, rects = [], [], []
+        for term, obs in per.items():
+            if term not in terms:
+                continue
+            vlines, term_rects = C.pvalue_regions(obs, spec.direction)
+            for x, is_dashed in vlines:
+                (dashed if is_dashed else solid).append({"term": term, "x": x})
+            for xmin, xmax in term_rects:
+                rects.append({"term": term, "xmin": xmin, "xmax": xmax})
+        if solid:
+            layers.append(_facet_vlines(solid, C._OBS_COLOR, dashed=False))
+        if dashed:
+            layers.append(_facet_vlines(dashed, C._OBS_COLOR, dashed=True))
+        if rects:
+            layers.append(_facet_rect(rects, C._OBS_COLOR))
+    else:
+        color = spec.color or C._SHADE_COLOR
+        rects, edges = [], []
+        for term, (lower, upper) in per.items():
+            if term not in terms:
+                continue
+            rects.append({"term": term, "xmin": lower, "xmax": upper})
+            edges.append({"term": term, "x": lower})
+            edges.append({"term": term, "x": upper})
+        if rects:
+            layers.append(_facet_rect(rects, color))
+        if edges:
+            layers.append(_facet_vlines(edges, color, dashed=False))
+    return gg + layers