Skip to content

Commit b9894fd

Browse files
lucasjamarclaude
lucasjamar
and
claude
committed
Support marginal_x/marginal_y="heatmap" in density_heatmap/density_contour
Draws a single-row/column heatmap strip in the margin, colored by the same z/histfunc aggregate as the main plot, sharing its coloraxis. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent de4f21b commit b9894fd

4 files changed

Lines changed: 149 additions & 6 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@ This project adheres to [Semantic Versioning](http://semver.org/).
44

55
## Unreleased
66

7+
### Added
8+
- Support `marginal_x`/`marginal_y="heatmap"` in `density_heatmap`/`density_contour`, drawing a single-row/column heatmap strip in the margin colored by the same `z`/`histfunc` aggregate as the main plot [[#5706](https://github.com/plotly/plotly.py/issues/5706)]
9+
710
### Fixed
811
- Fix `hex_to_rgb` parsing of 3-digit shorthand hexadecimal colors such as `#FFF` [[#5662](https://github.com/plotly/plotly.py/pull/5662)], with thanks to @genrichez for the contribution!
912
- Add `<!doctype html>` to the `to_html()` template to comply with modern web standards [[#5693](https://github.com/plotly/plotly.py/pull/5693)], with thanks to @mishrakushal for the contribution!

‎plotly/express/_chart_types.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,14 @@ def density_contour(
142142
"For `density_heatmap` and `density_contour` these values are used as the inputs to `histfunc`.",
143143
],
144144
histfunc=["The arguments to this function are the values of `z`."],
145+
marginal_x=[
146+
"Also supports `'heatmap'`, showing a single-row heatmap colored by the aggregate value. "
147+
"Uses a default colorscale since `density_contour` has no `color_continuous_scale` argument.",
148+
],
149+
marginal_y=[
150+
"Also supports `'heatmap'`, showing a single-column heatmap colored by the aggregate value. "
151+
"Uses a default colorscale since `density_contour` has no `color_continuous_scale` argument.",
152+
],
145153
),
146154
)
147155

@@ -214,6 +222,12 @@ def density_heatmap(
214222
histfunc=[
215223
"The arguments to this function are the values of `z`.",
216224
],
225+
marginal_x=[
226+
"Also supports `'heatmap'`, showing a single-row heatmap colored by the aggregate value.",
227+
],
228+
marginal_y=[
229+
"Also supports `'heatmap'`, showing a single-column heatmap colored by the aggregate value.",
230+
],
217231
),
218232
)
219233

‎plotly/express/_core.py‎

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -971,12 +971,48 @@ def make_trace_spec(args, constructor, attrs, trace_patch):
971971
),
972972
marginal=letter,
973973
)
974+
elif args["marginal_" + letter] == "heatmap":
975+
if constructor not in [go.Histogram2d, go.Histogram2dContour]:
976+
raise ValueError(
977+
"`marginal_x`/`marginal_y` value `'heatmap'` is only supported "
978+
"for `density_heatmap` and `density_contour`."
979+
)
980+
other_letter = "y" if letter == "x" else "x"
981+
heatmap_trace_patch = dict(
982+
coloraxis="coloraxis1", histfunc=args.get("histfunc"), **axis_map
983+
)
984+
# `nbinsx`/`nbinsy` are only a target bin count -- plotly.js's "nice
985+
# number" bin-sizing can still round to more than one bin. Force
986+
# exactly one bin by setting explicit bin edges covering the data.
987+
other_col = args["data_frame"].get_column(args[other_letter])
988+
other_min = nw.to_py_scalar(other_col.min())
989+
other_max = nw.to_py_scalar(other_col.max())
990+
span = (other_max - other_min) or 1
991+
pad = span * 0.001
992+
other_bins = dict(
993+
start=other_min - pad, end=other_max + pad, size=span + 2 * pad
994+
)
995+
if letter == "x":
996+
heatmap_trace_patch["xbingroup"] = "x"
997+
heatmap_trace_patch["ybins"] = other_bins
998+
else:
999+
heatmap_trace_patch["ybingroup"] = "y"
1000+
heatmap_trace_patch["xbins"] = other_bins
1001+
trace_spec = TraceSpec(
1002+
constructor=go.Histogram2d,
1003+
attrs=[letter, other_letter, "z"],
1004+
trace_patch=heatmap_trace_patch,
1005+
marginal=letter,
1006+
)
9741007
else:
9751008
raise ValueError(
9761009
f"Invalid value '{args['marginal_' + letter]}' for `marginal_{letter}`. "
977-
"Supported marginal plot types are: 'rug', 'box', 'violin', 'histogram'."
1010+
"Supported marginal plot types are: "
1011+
"'rug', 'box', 'violin', 'histogram', 'heatmap'."
9781012
)
979-
if "color" in attrs or "color" not in args:
1013+
if trace_spec.constructor != go.Histogram2d and (
1014+
"color" in attrs or "color" not in args
1015+
):
9801016
if "marker" not in trace_spec.trace_patch:
9811017
trace_spec.trace_patch["marker"] = dict()
9821018
first_default_color = args["color_continuous_scale"][0]
@@ -2337,6 +2373,10 @@ def infer_config(args, constructor, trace_patch, layout_patch):
23372373
if constructor in [go.Histogram2d, go.Densitymap, go.Densitymapbox]:
23382374
show_colorbar = True
23392375
trace_patch["coloraxis"] = "coloraxis1"
2376+
elif constructor == go.Histogram2dContour and (
2377+
args.get("marginal_x") == "heatmap" or args.get("marginal_y") == "heatmap"
2378+
):
2379+
show_colorbar = True
23402380

23412381
if "opacity" in args:
23422382
if args["opacity"] is None:
@@ -2630,6 +2670,10 @@ def make_figure(args, constructor, trace_patch=None, layout_patch=None):
26302670
trace_spec.constructor in [go.Histogram]
26312671
and m.variable in ["symbol", "dash"]
26322672
)
2673+
or (
2674+
trace_spec.constructor == go.Histogram2d
2675+
and m.variable in ["symbol", "pattern", "dash", "color"]
2676+
)
26332677
):
26342678
pass
26352679
elif (
@@ -2727,17 +2771,18 @@ def make_figure(args, constructor, trace_patch=None, layout_patch=None):
27272771
if show_colorbar:
27282772
colorvar = (
27292773
"z"
2730-
if constructor in [go.Histogram2d, go.Densitymap, go.Densitymapbox]
2774+
if constructor
2775+
in [go.Histogram2d, go.Histogram2dContour, go.Densitymap, go.Densitymapbox]
27312776
else "color"
27322777
)
2733-
range_color = args["range_color"] or [None, None]
2778+
range_color = args.get("range_color") or [None, None]
27342779

27352780
colorscale_validator = ColorscaleValidator("colorscale", "make_figure")
27362781
coloraxis_dict = dict(
27372782
colorscale=colorscale_validator.validate_coerce(
2738-
args["color_continuous_scale"]
2783+
args.get("color_continuous_scale")
27392784
),
2740-
cmid=args["color_continuous_midpoint"],
2785+
cmid=args.get("color_continuous_midpoint"),
27412786
cmin=range_color[0],
27422787
cmax=range_color[1],
27432788
colorbar=dict(

‎tests/test_optional/test_px/test_marginals.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,87 @@ def test_single_marginals(backend, px_fn, marginal, orientation):
2626
assert len(fig.data) == 1 + (marginal is not None)
2727

2828

29+
@pytest.mark.parametrize("px_fn", [px.density_heatmap, px.density_contour])
30+
def test_marginal_heatmap_uses_z_and_histfunc(backend, px_fn):
31+
df = px.data.tips(return_type=backend)
32+
# backend-independent reference for min/max, since e.g. pyarrow columns don't
33+
# support .min()/.max() directly
34+
pdf = px.data.tips()
35+
36+
fig = px_fn(
37+
df,
38+
x="total_bill",
39+
y="tip",
40+
z="size",
41+
histfunc="sum",
42+
marginal_x="heatmap",
43+
marginal_y="heatmap",
44+
)
45+
assert len(fig.data) == 3
46+
marginal_x_trace, marginal_y_trace = fig.data[1], fig.data[2]
47+
48+
assert marginal_x_trace.type == "histogram2d"
49+
assert marginal_x_trace.coloraxis == "coloraxis"
50+
assert marginal_x_trace.histfunc == "sum"
51+
# a single bin covering the full y range, so the strip is exactly one row
52+
assert marginal_x_trace.ybins.start <= pdf["tip"].min()
53+
assert marginal_x_trace.ybins.end >= pdf["tip"].max()
54+
assert marginal_x_trace.ybins.size >= pdf["tip"].max() - pdf["tip"].min()
55+
56+
assert marginal_y_trace.type == "histogram2d"
57+
assert marginal_y_trace.coloraxis == "coloraxis"
58+
assert marginal_y_trace.histfunc == "sum"
59+
# a single bin covering the full x range, so the strip is exactly one column
60+
assert marginal_y_trace.xbins.start <= pdf["total_bill"].min()
61+
assert marginal_y_trace.xbins.end >= pdf["total_bill"].max()
62+
assert marginal_y_trace.xbins.size >= pdf["total_bill"].max() - pdf["total_bill"].min()
63+
64+
assert fig.layout.coloraxis.colorbar.title.text == "sum of size"
65+
66+
67+
@pytest.mark.parametrize("px_fn", [px.density_heatmap, px.density_contour])
68+
def test_marginal_heatmap_without_z(backend, px_fn):
69+
df = px.data.tips(return_type=backend)
70+
71+
fig = px_fn(
72+
df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap"
73+
)
74+
marginal_x_trace, marginal_y_trace = fig.data[1], fig.data[2]
75+
76+
assert marginal_x_trace.type == "histogram2d"
77+
assert marginal_x_trace.coloraxis == "coloraxis"
78+
assert marginal_x_trace.histfunc is None
79+
80+
assert marginal_y_trace.type == "histogram2d"
81+
assert marginal_y_trace.coloraxis == "coloraxis"
82+
assert marginal_y_trace.histfunc is None
83+
84+
assert fig.layout.coloraxis.colorbar.title.text == "count"
85+
86+
87+
def test_marginal_heatmap_unsupported_chart_type_raises():
88+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
89+
px.scatter(x=[1, 2, 3], y=[2, 3, 4], marginal_x="heatmap")
90+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
91+
px.scatter(x=[1, 2, 3], y=[2, 3, 4], marginal_y="heatmap")
92+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
93+
px.histogram(x=[1, 2, 3], marginal="heatmap")
94+
95+
96+
def test_marginal_heatmap_with_discrete_color(backend): # density_contour + line.color
97+
df = px.data.tips(return_type=backend)
98+
fig = px.density_contour(
99+
df, x="total_bill", y="tip", color="sex", marginal_x="heatmap"
100+
)
101+
assert len(fig.data) == 4
102+
assert [t.type for t in fig.data] == [
103+
"histogram2dcontour",
104+
"histogram2d",
105+
"histogram2dcontour",
106+
"histogram2d",
107+
]
108+
109+
29110
def test_unsupported_marginal_raises_clear_error(): # issue 4654
30111
# An unsupported marginal type used to fail deep inside make_figure with a
31112
# cryptic "'NoneType' object has no attribute 'constructor'". It should

0 commit comments

Comments
 (0)