Support marginal_x/marginal_y="heatmap" in density_heatmap/density_contour

CHANGELOG.md

@@ -4,6 +4,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 
+### Added
+- 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!
 
 ## [7.0.0] - 2026-08-25
 

doc/python/marginal-plots.md

@@ -39,7 +39,7 @@ Marginal distribution plots are small subplots above or to the right of a main p
 
 ### Scatter Plot Marginals
 
-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). 
+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).
 
 Marginal plots are linked to the main plot: try zooming or panning on the main plot.
 
@@ -59,6 +59,27 @@ fig = px.density_heatmap(df, x="sepal_length", y="sepal_width", marginal_x="box"
 fig.show()
 ```
 
+### Marginal Heatmaps on Density Heatmaps
+
+`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.
+
+```python
+import plotly.express as px
+df = px.data.tips()
+fig = px.density_heatmap(df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap")
+fig.show()
+```
+
+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:
+
+```python
+import plotly.express as px
+df = px.data.tips()
+fig = px.density_heatmap(df, x="total_bill", y="tip", z="size", histfunc="avg",
+                          marginal_x="heatmap", marginal_y="heatmap", text_auto=".1f")
+fig.show()
+```
+
 ### Marginal Plots and Color
 
 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

@@ -214,6 +214,15 @@ def density_heatmap(
         histfunc=[
             "The arguments to this function are the values of `z`.",
         ],
+        marginal_x=[
+            "Also supports `'heatmap'`, showing a single-row heatmap colored by the aggregate value.",
+        ],
+        marginal_y=[
+            "Also supports `'heatmap'`, showing a single-column heatmap colored by the aggregate value.",
+        ],
+        text_auto=[
+            "Also applies to `marginal_x`/`marginal_y='heatmap'`, in which case the z values are always displayed.",
+        ],
     ),
 )
 

plotly/express/_core.py

@@ -944,12 +944,55 @@ def make_trace_spec(args, constructor, attrs, trace_patch):
                     ),
                     marginal=letter,
                 )
+            elif args["marginal_" + letter] == "heatmap":
+                if constructor != go.Histogram2d:
+                    raise ValueError(
+                        "`marginal_x`/`marginal_y` value `'heatmap'` is only supported "
+                        "for `density_heatmap`."
+                    )
+                other_letter = "y" if letter == "x" else "x"
+                heatmap_trace_patch = dict(
+                    coloraxis="coloraxis1", histfunc=args.get("histfunc"), **axis_map
+                )
+                # `nbinsx`/`nbinsy` are only a target bin count -- plotly.js's "nice
+                # number" bin-sizing can still round to more than one bin. Force
+                # exactly one bin by setting explicit bin edges covering the data.
+                other_col = args["data_frame"].get_column(args[other_letter])
+                other_min = nw.to_py_scalar(other_col.min())
+                other_max = nw.to_py_scalar(other_col.max())
+                span = (other_max - other_min) or 1
+                pad = span * 0.001
+                other_bins = dict(
+                    start=other_min - pad, end=other_max + pad, size=span + 2 * pad
+                )
+                if letter == "x":
+                    heatmap_trace_patch["xbingroup"] = "x"
+                    heatmap_trace_patch["ybins"] = other_bins
+                else:
+                    heatmap_trace_patch["ybingroup"] = "y"
+                    heatmap_trace_patch["xbins"] = other_bins
+                if args.get("text_auto", False) is not False:
+                    if args["text_auto"] is True:
+                        heatmap_trace_patch["texttemplate"] = "%{z}"
+                    else:
+                        heatmap_trace_patch["texttemplate"] = (
+                            "%{z:" + args["text_auto"] + "}"
+                        )
+                trace_spec = TraceSpec(
+                    constructor=go.Histogram2d,
+                    attrs=[letter, other_letter, "z"],
+                    trace_patch=heatmap_trace_patch,
+                    marginal=letter,
+                )
             else:
                 raise ValueError(
                     f"Invalid value '{args['marginal_' + letter]}' for `marginal_{letter}`. "
-                    "Supported marginal plot types are: 'rug', 'box', 'violin', 'histogram'."
+                    "Supported marginal plot types are: "
+                    "'rug', 'box', 'violin', 'histogram', 'heatmap'."
                 )
-            if "color" in attrs or "color" not in args:
+            if trace_spec.constructor != go.Histogram2d and (
+                "color" in attrs or "color" not in args
+            ):
                 if "marker" not in trace_spec.trace_patch:
                     trace_spec.trace_patch["marker"] = dict()
                 first_default_color = args["color_continuous_scale"][0]

tests/test_optional/test_px/test_marginals.py

@@ -26,6 +26,126 @@ def test_single_marginals(backend, px_fn, marginal, orientation):
     assert len(fig.data) == 1 + (marginal is not None)
 
 
+def test_marginal_heatmap_uses_z_and_histfunc(backend):
+    df = px.data.tips(return_type=backend)
+    # backend-independent reference for min/max, since e.g. pyarrow columns don't
+    # support .min()/.max() directly
+    pdf = px.data.tips()
+
+    fig = px.density_heatmap(
+        df,
+        x="total_bill",
+        y="tip",
+        z="size",
+        histfunc="sum",
+        marginal_x="heatmap",
+        marginal_y="heatmap",
+    )
+    assert len(fig.data) == 3
+    main_trace, marginal_x_trace, marginal_y_trace = (
+        fig.data[0],
+        fig.data[1],
+        fig.data[2],
+    )
+
+    assert marginal_x_trace.type == "histogram2d"
+    assert marginal_x_trace.coloraxis == "coloraxis"
+    assert marginal_x_trace.histfunc == "sum"
+    # a single bin covering the full y range, so the strip is exactly one row
+    assert marginal_x_trace.ybins.start <= pdf["tip"].min()
+    assert marginal_x_trace.ybins.end >= pdf["tip"].max()
+    assert marginal_x_trace.ybins.size >= pdf["tip"].max() - pdf["tip"].min()
+
+    assert marginal_y_trace.type == "histogram2d"
+    assert marginal_y_trace.coloraxis == "coloraxis"
+    assert marginal_y_trace.histfunc == "sum"
+    # a single bin covering the full x range, so the strip is exactly one column
+    assert marginal_y_trace.xbins.start <= pdf["total_bill"].min()
+    assert marginal_y_trace.xbins.end >= pdf["total_bill"].max()
+    assert (
+        marginal_y_trace.xbins.size >= pdf["total_bill"].max() - pdf["total_bill"].min()
+    )
+
+    assert fig.layout.coloraxis.colorbar.title.text == "sum of size"
+
+    # Ensure the x, y, and z data for the marginal heatmaps are consistent with the main heatmap
+    assert (marginal_x_trace.x == main_trace.x).all()
+    assert (marginal_x_trace.y == main_trace.y).all()
+    assert (marginal_x_trace.z == main_trace.z).all()
+    assert (marginal_y_trace.x == main_trace.x).all()
+    assert (marginal_y_trace.y == main_trace.y).all()
+    assert (marginal_y_trace.z == main_trace.z).all()
+
+
+def test_marginal_heatmap_without_z(backend):
+    df = px.data.tips(return_type=backend)
+
+    fig = px.density_heatmap(
+        df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap"
+    )
+    main_trace, marginal_x_trace, marginal_y_trace = (
+        fig.data[0],
+        fig.data[1],
+        fig.data[2],
+    )
+
+    assert marginal_x_trace.type == "histogram2d"
+    assert marginal_x_trace.coloraxis == "coloraxis"
+    assert marginal_x_trace.histfunc is None
+
+    assert marginal_y_trace.type == "histogram2d"
+    assert marginal_y_trace.coloraxis == "coloraxis"
+    assert marginal_y_trace.histfunc is None
+
+    assert fig.layout.coloraxis.colorbar.title.text == "count"
+
+    # Ensure the x and y data for the marginal heatmaps are consistent with the main heatmap
+    assert (marginal_x_trace.x == main_trace.x).all()
+    assert (marginal_x_trace.y == main_trace.y).all()
+    assert (marginal_y_trace.x == main_trace.x).all()
+    assert (marginal_y_trace.y == main_trace.y).all()
+
+
+@pytest.mark.parametrize("text_auto", [True, ".1f"])
+def test_marginal_heatmap_text_auto(backend, text_auto):
+    df = px.data.tips(return_type=backend)
+
+    fig = px.density_heatmap(
+        df,
+        x="total_bill",
+        y="tip",
+        marginal_x="heatmap",
+        marginal_y="heatmap",
+        text_auto=text_auto,
+    )
+    expected = "%{z}" if text_auto is True else "%{z:" + text_auto + "}"
+    for trace in fig.data:
+        assert trace.texttemplate == expected
+
+
+def test_marginal_heatmap_no_text_auto(backend):
+    df = px.data.tips(return_type=backend)
+
+    fig = px.density_heatmap(
+        df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap"
+    )
+    for trace in fig.data:
+        assert trace.texttemplate is None
+
+
+def test_marginal_heatmap_unsupported_chart_type_raises():
+    with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
+        px.scatter(x=[1, 2, 3], y=[2, 3, 4], marginal_x="heatmap")
+    with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
+        px.scatter(x=[1, 2, 3], y=[2, 3, 4], marginal_y="heatmap")
+    with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
+        px.histogram(x=[1, 2, 3], marginal="heatmap")
+    with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
+        px.density_contour(x=[1, 2, 3], y=[2, 3, 4], marginal_x="heatmap")
+    with pytest.raises(ValueError, match="only supported for `density_heatmap`"):
+        px.density_contour(x=[1, 2, 3], y=[2, 3, 4], marginal_y="heatmap")
+
+
 @pytest.mark.parametrize("px_fn", [px.density_heatmap, px.density_contour])
 def test_marginal_histogram_uses_z_and_histfunc(backend, px_fn):  # issue 3521
     df = px.data.tips(return_type=backend)