diff --git a/CHANGELOG.md b/CHANGELOG.md index 45ee3cc..3af8027 100644 --- a/CHANGELOG.md +++ b/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"`. diff --git a/moderndive/infer/core.py b/moderndive/infer/core.py index 2b7a1f3..92502f6 100644 --- a/moderndive/infer/core.py +++ b/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( diff --git a/moderndive/infer/viz/__init__.py b/moderndive/infer/viz/__init__.py index 7963b31..8b2c2ce 100644 --- a/moderndive/infer/viz/__init__.py +++ b/moderndive/infer/viz/__init__.py @@ -40,7 +40,14 @@ @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 @@ -48,6 +55,7 @@ class ShadeSpec: 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) diff --git a/moderndive/infer/viz/_plotly.py b/moderndive/infer/viz/_plotly.py index 0b098d1..8cca4bc 100644 --- a/moderndive/infer/viz/_plotly.py +++ b/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 diff --git a/moderndive/infer/viz/_plotnine.py b/moderndive/infer/viz/_plotnine.py index 7fa3cd0..3f3c268 100644 --- a/moderndive/infer/viz/_plotnine.py +++ b/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 diff --git a/tests/test_viz.py b/tests/test_viz.py index 5b08e17..364ddfd 100644 --- a/tests/test_viz.py +++ b/tests/test_viz.py @@ -18,7 +18,21 @@ specify, visualize, ) -from moderndive.infer.viz import InferPlot, ShadeSpec +from moderndive.infer.viz import InferPlot, ShadeSpec, visualize_fit + + +def _fit_obs_null_ci(): + sar = md.load_saratoga_houses() + f = "price ~ living_area + bedrooms" + obs = specify(sar, formula=f).fit() + dist = specify(sar, formula=f).generate(reps=60, type="bootstrap", seed=1).fit() + null = ( + specify(sar, formula=f) + .hypothesize(null="independence") + .generate(reps=60, type="permute", seed=1) + .fit() + ) + return obs, dist, null, dist.get_confidence_interval(level=0.95) def _boot(): @@ -192,3 +206,85 @@ def test_apply_shade_on_empty_plotly_figure(): empty = InferPlot(go.Figure(), "plotly") shaded = empty + shade_confidence_interval((1.0, 2.0)) assert isinstance(shaded.figure, go.Figure) + + +# ---- per-facet shading for faceted regression-fit plots (both engines) ---- + + +@pytest.mark.parametrize("direction", ["right", "left", "two-sided"]) +def test_fit_per_facet_pvalue_plotly(direction): + obs, _dist, null, _ci = _fit_obs_null_ci() + p = visualize_fit(null, engine="plotly") + shade_p_value(obs_stat=obs, direction=direction) + assert p.terms == ["intercept", "living_area", "bedrooms"] + # one shaded region per term (plus vlines); shapes scale with the 3 facets + assert len(p.figure.layout.shapes) >= 3 * 3 if direction == "two-sided" else 3 * 2 + + +def test_fit_per_facet_ci_plotly(): + _obs, dist, _null, ci = _fit_obs_null_ci() + p = visualize_fit(dist, engine="plotly") + shade_confidence_interval(ci) + # 3 terms x (1 vrect + 2 vlines) = 9 shapes + assert len(p.figure.layout.shapes) == 9 + + +def test_fit_per_facet_pvalue_plotnine(): + obs, _dist, null, _ci = _fit_obs_null_ci() + base = visualize_fit(null, engine="plotnine") + shaded = base + shade_p_value(obs_stat=obs, direction="two-sided") + assert len(shaded.gg.layers) > len(base.gg.layers) + assert isinstance(shaded.gg, ggplot) + + +def test_fit_per_facet_ci_plotnine_keyword(): + _obs, dist, _null, ci = _fit_obs_null_ci() + # keyword form on FitResult.visualize, custom color + g = dist.visualize(engine="plotnine", shade_ci=ci) + base = visualize_fit(dist, engine="plotnine") + assert len(g.gg.layers) > len(base.gg.layers) + + +def test_fit_shade_pvalue_keyword_dict_form(): + obs, _dist, null, _ci = _fit_obs_null_ci() + p = null.visualize(engine="plotly", shade_pvalue={"obs_stat": obs, "direction": "right"}) + assert len(p.figure.layout.shapes) >= 3 + + +def test_per_term_obs_accepts_table_and_dict(): + obs, _dist, _null, _ci = _fit_obs_null_ci() + # term/estimate frame, term/stat frame, and dict all yield per-term specs + spec_df = shade_p_value(obs.data, direction="right") + spec_stat = shade_p_value(obs.data.rename({"estimate": "stat"}), direction="right") + spec_dict = shade_p_value({"living_area": 90.0, "bedrooms": -7000.0}, direction="right") + for spec in (spec_df, spec_stat, spec_dict): + assert spec.per_term is not None and spec.obs_stat is None + + +def test_per_term_shade_on_nonfacet_raises(): + obs, _dist, _null, _ci = _fit_obs_null_ci() + boot = _boot() + with pytest.raises(TypeError): + visualize(boot, engine="plotly") + shade_p_value(obs_stat=obs, direction="right") + + +def test_per_facet_skips_unknown_terms(): + import polars as pl + + _obs, dist, _null, ci = _fit_obs_null_ci() + # add a term not present in the facets; it must be silently skipped (both engines) + extra = pl.concat([ci, ci.head(1).with_columns(pl.lit("ghost").alias("term"))]) + p = visualize_fit(dist, engine="plotly") + shade_confidence_interval(extra) + assert len(p.figure.layout.shapes) == 9 # ghost term ignored + g = visualize_fit(dist, engine="plotnine") + shade_confidence_interval(extra) + assert isinstance(g.gg, ggplot) + # p-value path: a dict with a term not among the facets is skipped too + obs_with_ghost = {"living_area": 90.0, "ghost": 1.0} + pp = visualize_fit(dist, engine="plotly") + shade_p_value(obs_with_ghost, direction="right") + assert len(pp.figure.layout.shapes) == 2 # only living_area shaded (1 vline + 1 rect) + gp = visualize_fit(dist, engine="plotnine") + shade_p_value(obs_with_ghost, direction="right") + assert isinstance(gp.gg, ggplot) + + +def test_subplot_xrange_fallback(): + from moderndive.infer.viz._plotly import _subplot_xrange + + assert _subplot_xrange(go.Figure(), 1) == (-1.0, 1.0)