Skip to content

Commit 15f33d2

Browse files
authored
Merge pull request #5707 from lucasjamar/marginal-heatmap-feature
Support marginal_x/marginal_y="heatmap" in density_heatmap/density_contour
2 parents ca01aba + 996b31b commit 15f33d2

5 files changed

Lines changed: 198 additions & 3 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ 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`, drawing a single-row/column heatmap strip in the margin colored by the same `z`/`histfunc` aggregate as the main plot and sharing its color scale [[#5706](https://github.com/plotly/plotly.py/issues/5706)], with thanks to @lucasjamar for the contribution!
79

810
## [7.0.0] - 2026-08-25
911

‎doc/python/marginal-plots.md‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ Marginal distribution plots are small subplots above or to the right of a main p
3939

4040
### Scatter Plot Marginals
4141

42-
The `marginal_x` and `marginal_y` arguments accept one of `"histogram"`, `"rug"`, `"box"`, or `"violin"` (see also how to create [histograms](/python/histograms/), [box plots](/python/box-plots/) and [violin plots](/python/violin-plots/) as the main figure).
42+
The `marginal_x` and `marginal_y` arguments accept one of `"histogram"`, `"rug"`, `"box"`, or `"violin"` (see also how to create [histograms](/python/histograms/), [box plots](/python/box-plots/) and [violin plots](/python/violin-plots/) as the main figure), plus `"heatmap"` for `density_heatmap` (see below).
4343

4444
Marginal plots are linked to the main plot: try zooming or panning on the main plot.
4545

@@ -59,6 +59,27 @@ fig = px.density_heatmap(df, x="sepal_length", y="sepal_width", marginal_x="box"
5959
fig.show()
6060
```
6161

62+
### Marginal Heatmaps on Density Heatmaps
63+
64+
`marginal_x` and `marginal_y` also accept `"heatmap"` for [`density_heatmap`](/python/2D-Histogram/). This draws a single-row or single-column heatmap strip, colored by the same aggregate (`histfunc` of `z`, or count by default) as the main plot, and sharing its color scale. This is not supported for `density_contour`, since a contour plot's colorbar is discrete and not compatible with the continuous colorbar used by the marginal heatmap.
65+
66+
```python
67+
import plotly.express as px
68+
df = px.data.tips()
69+
fig = px.density_heatmap(df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap")
70+
fig.show()
71+
```
72+
73+
Set `text_auto=True` to display the aggregate value as text on both the main plot and the marginal heatmap strips, or pass a [d3-format](https://github.com/d3/d3-format) string such as `".2f"` to control the number of decimal places:
74+
75+
```python
76+
import plotly.express as px
77+
df = px.data.tips()
78+
fig = px.density_heatmap(df, x="total_bill", y="tip", z="size", histfunc="avg",
79+
marginal_x="heatmap", marginal_y="heatmap", text_auto=".1f")
80+
fig.show()
81+
```
82+
6283
### Marginal Plots and Color
6384

6485
Marginal plots respect the `color` argument as well, and are linked to the respective legend elements. Try clicking on the legend items.

‎plotly/express/_chart_types.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,15 @@ def density_heatmap(
214214
histfunc=[
215215
"The arguments to this function are the values of `z`.",
216216
],
217+
marginal_x=[
218+
"Also supports `'heatmap'`, showing a single-row heatmap colored by the aggregate value.",
219+
],
220+
marginal_y=[
221+
"Also supports `'heatmap'`, showing a single-column heatmap colored by the aggregate value.",
222+
],
223+
text_auto=[
224+
"Also applies to `marginal_x`/`marginal_y='heatmap'`, in which case the z values are always displayed.",
225+
],
217226
),
218227
)
219228

‎plotly/express/_core.py‎

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -944,12 +944,55 @@ def make_trace_spec(args, constructor, attrs, trace_patch):
944944
),
945945
marginal=letter,
946946
)
947+
elif args["marginal_" + letter] == "heatmap":
948+
if constructor != go.Histogram2d:
949+
raise ValueError(
950+
"`marginal_x`/`marginal_y` value `'heatmap'` is only supported "
951+
"for `density_heatmap`."
952+
)
953+
other_letter = "y" if letter == "x" else "x"
954+
heatmap_trace_patch = dict(
955+
coloraxis="coloraxis1", histfunc=args.get("histfunc"), **axis_map
956+
)
957+
# `nbinsx`/`nbinsy` are only a target bin count -- plotly.js's "nice
958+
# number" bin-sizing can still round to more than one bin. Force
959+
# exactly one bin by setting explicit bin edges covering the data.
960+
other_col = args["data_frame"].get_column(args[other_letter])
961+
other_min = nw.to_py_scalar(other_col.min())
962+
other_max = nw.to_py_scalar(other_col.max())
963+
span = (other_max - other_min) or 1
964+
pad = span * 0.001
965+
other_bins = dict(
966+
start=other_min - pad, end=other_max + pad, size=span + 2 * pad
967+
)
968+
if letter == "x":
969+
heatmap_trace_patch["xbingroup"] = "x"
970+
heatmap_trace_patch["ybins"] = other_bins
971+
else:
972+
heatmap_trace_patch["ybingroup"] = "y"
973+
heatmap_trace_patch["xbins"] = other_bins
974+
if args.get("text_auto", False) is not False:
975+
if args["text_auto"] is True:
976+
heatmap_trace_patch["texttemplate"] = "%{z}"
977+
else:
978+
heatmap_trace_patch["texttemplate"] = (
979+
"%{z:" + args["text_auto"] + "}"
980+
)
981+
trace_spec = TraceSpec(
982+
constructor=go.Histogram2d,
983+
attrs=[letter, other_letter, "z"],
984+
trace_patch=heatmap_trace_patch,
985+
marginal=letter,
986+
)
947987
else:
948988
raise ValueError(
949989
f"Invalid value '{args['marginal_' + letter]}' for `marginal_{letter}`. "
950-
"Supported marginal plot types are: 'rug', 'box', 'violin', 'histogram'."
990+
"Supported marginal plot types are: "
991+
"'rug', 'box', 'violin', 'histogram', 'heatmap'."
951992
)
952-
if "color" in attrs or "color" not in args:
993+
if trace_spec.constructor != go.Histogram2d and (
994+
"color" in attrs or "color" not in args
995+
):
953996
if "marker" not in trace_spec.trace_patch:
954997
trace_spec.trace_patch["marker"] = dict()
955998
first_default_color = args["color_continuous_scale"][0]

‎tests/test_optional/test_px/test_marginals.py‎

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

2828

29+
def test_marginal_heatmap_uses_z_and_histfunc(backend):
30+
df = px.data.tips(return_type=backend)
31+
# backend-independent reference for min/max, since e.g. pyarrow columns don't
32+
# support .min()/.max() directly
33+
pdf = px.data.tips()
34+
35+
fig = px.density_heatmap(
36+
df,
37+
x="total_bill",
38+
y="tip",
39+
z="size",
40+
histfunc="sum",
41+
marginal_x="heatmap",
42+
marginal_y="heatmap",
43+
)
44+
assert len(fig.data) == 3
45+
main_trace, marginal_x_trace, marginal_y_trace = (
46+
fig.data[0],
47+
fig.data[1],
48+
fig.data[2],
49+
)
50+
51+
assert marginal_x_trace.type == "histogram2d"
52+
assert marginal_x_trace.coloraxis == "coloraxis"
53+
assert marginal_x_trace.histfunc == "sum"
54+
# a single bin covering the full y range, so the strip is exactly one row
55+
assert marginal_x_trace.ybins.start <= pdf["tip"].min()
56+
assert marginal_x_trace.ybins.end >= pdf["tip"].max()
57+
assert marginal_x_trace.ybins.size >= pdf["tip"].max() - pdf["tip"].min()
58+
59+
assert marginal_y_trace.type == "histogram2d"
60+
assert marginal_y_trace.coloraxis == "coloraxis"
61+
assert marginal_y_trace.histfunc == "sum"
62+
# a single bin covering the full x range, so the strip is exactly one column
63+
assert marginal_y_trace.xbins.start <= pdf["total_bill"].min()
64+
assert marginal_y_trace.xbins.end >= pdf["total_bill"].max()
65+
assert (
66+
marginal_y_trace.xbins.size >= pdf["total_bill"].max() - pdf["total_bill"].min()
67+
)
68+
69+
assert fig.layout.coloraxis.colorbar.title.text == "sum of size"
70+
71+
# Ensure the x, y, and z data for the marginal heatmaps are consistent with the main heatmap
72+
assert (marginal_x_trace.x == main_trace.x).all()
73+
assert (marginal_x_trace.y == main_trace.y).all()
74+
assert (marginal_x_trace.z == main_trace.z).all()
75+
assert (marginal_y_trace.x == main_trace.x).all()
76+
assert (marginal_y_trace.y == main_trace.y).all()
77+
assert (marginal_y_trace.z == main_trace.z).all()
78+
79+
80+
def test_marginal_heatmap_without_z(backend):
81+
df = px.data.tips(return_type=backend)
82+
83+
fig = px.density_heatmap(
84+
df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap"
85+
)
86+
main_trace, marginal_x_trace, marginal_y_trace = (
87+
fig.data[0],
88+
fig.data[1],
89+
fig.data[2],
90+
)
91+
92+
assert marginal_x_trace.type == "histogram2d"
93+
assert marginal_x_trace.coloraxis == "coloraxis"
94+
assert marginal_x_trace.histfunc is None
95+
96+
assert marginal_y_trace.type == "histogram2d"
97+
assert marginal_y_trace.coloraxis == "coloraxis"
98+
assert marginal_y_trace.histfunc is None
99+
100+
assert fig.layout.coloraxis.colorbar.title.text == "count"
101+
102+
# Ensure the x and y data for the marginal heatmaps are consistent with the main heatmap
103+
assert (marginal_x_trace.x == main_trace.x).all()
104+
assert (marginal_x_trace.y == main_trace.y).all()
105+
assert (marginal_y_trace.x == main_trace.x).all()
106+
assert (marginal_y_trace.y == main_trace.y).all()
107+
108+
109+
@pytest.mark.parametrize("text_auto", [True, ".1f"])
110+
def test_marginal_heatmap_text_auto(backend, text_auto):
111+
df = px.data.tips(return_type=backend)
112+
113+
fig = px.density_heatmap(
114+
df,
115+
x="total_bill",
116+
y="tip",
117+
marginal_x="heatmap",
118+
marginal_y="heatmap",
119+
text_auto=text_auto,
120+
)
121+
expected = "%{z}" if text_auto is True else "%{z:" + text_auto + "}"
122+
for trace in fig.data:
123+
assert trace.texttemplate == expected
124+
125+
126+
def test_marginal_heatmap_no_text_auto(backend):
127+
df = px.data.tips(return_type=backend)
128+
129+
fig = px.density_heatmap(
130+
df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap"
131+
)
132+
for trace in fig.data:
133+
assert trace.texttemplate is None
134+
135+
136+
def test_marginal_heatmap_unsupported_chart_type_raises():
137+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
138+
px.scatter(x=[1, 2, 3], y=[2, 3, 4], marginal_x="heatmap")
139+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
140+
px.scatter(x=[1, 2, 3], y=[2, 3, 4], marginal_y="heatmap")
141+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
142+
px.histogram(x=[1, 2, 3], marginal="heatmap")
143+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
144+
px.density_contour(x=[1, 2, 3], y=[2, 3, 4], marginal_x="heatmap")
145+
with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
146+
px.density_contour(x=[1, 2, 3], y=[2, 3, 4], marginal_y="heatmap")
147+
148+
29149
@pytest.mark.parametrize("px_fn", [px.density_heatmap, px.density_contour])
30150
def test_marginal_histogram_uses_z_and_histfunc(backend, px_fn): # issue 3521
31151
df = px.data.tips(return_type=backend)

0 commit comments

Comments
 (0)