diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f4cc3ac0e..df849d949b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ This project adheres to [Semantic Versioning](http://semver.org/). ## Unreleased +### Added +- 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)] + ### Fixed - 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! - Add `` 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! diff --git a/plotly/express/_chart_types.py b/plotly/express/_chart_types.py index 2b4b10184e..b3dd365263 100644 --- a/plotly/express/_chart_types.py +++ b/plotly/express/_chart_types.py @@ -142,6 +142,14 @@ def density_contour( "For `density_heatmap` and `density_contour` these values are used as the inputs to `histfunc`.", ], 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. " + "Uses a default colorscale since `density_contour` has no `color_continuous_scale` argument.", + ], + marginal_y=[ + "Also supports `'heatmap'`, showing a single-column heatmap colored by the aggregate value. " + "Uses a default colorscale since `density_contour` has no `color_continuous_scale` argument.", + ], ), ) @@ -214,6 +222,12 @@ 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.", + ], ), ) diff --git a/plotly/express/_core.py b/plotly/express/_core.py index 1b599cb43c..b89eee384d 100644 --- a/plotly/express/_core.py +++ b/plotly/express/_core.py @@ -971,12 +971,48 @@ def make_trace_spec(args, constructor, attrs, trace_patch): ), marginal=letter, ) + elif args["marginal_" + letter] == "heatmap": + if constructor not in [go.Histogram2d, go.Histogram2dContour]: + raise ValueError( + "`marginal_x`/`marginal_y` value `'heatmap'` is only supported " + "for `density_heatmap` and `density_contour`." + ) + 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 + 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] @@ -2337,6 +2373,10 @@ def infer_config(args, constructor, trace_patch, layout_patch): if constructor in [go.Histogram2d, go.Densitymap, go.Densitymapbox]: show_colorbar = True trace_patch["coloraxis"] = "coloraxis1" + elif constructor == go.Histogram2dContour and ( + args.get("marginal_x") == "heatmap" or args.get("marginal_y") == "heatmap" + ): + show_colorbar = True if "opacity" in args: if args["opacity"] is None: @@ -2630,6 +2670,10 @@ def make_figure(args, constructor, trace_patch=None, layout_patch=None): trace_spec.constructor in [go.Histogram] and m.variable in ["symbol", "dash"] ) + or ( + trace_spec.constructor == go.Histogram2d + and m.variable in ["symbol", "pattern", "dash", "color"] + ) ): pass elif ( @@ -2727,17 +2771,18 @@ def make_figure(args, constructor, trace_patch=None, layout_patch=None): if show_colorbar: colorvar = ( "z" - if constructor in [go.Histogram2d, go.Densitymap, go.Densitymapbox] + if constructor + in [go.Histogram2d, go.Histogram2dContour, go.Densitymap, go.Densitymapbox] else "color" ) - range_color = args["range_color"] or [None, None] + range_color = args.get("range_color") or [None, None] colorscale_validator = ColorscaleValidator("colorscale", "make_figure") coloraxis_dict = dict( colorscale=colorscale_validator.validate_coerce( - args["color_continuous_scale"] + args.get("color_continuous_scale") ), - cmid=args["color_continuous_midpoint"], + cmid=args.get("color_continuous_midpoint"), cmin=range_color[0], cmax=range_color[1], colorbar=dict( diff --git a/tests/test_optional/test_px/test_marginals.py b/tests/test_optional/test_px/test_marginals.py index 40d3019f2d..dc7e903626 100644 --- a/tests/test_optional/test_px/test_marginals.py +++ b/tests/test_optional/test_px/test_marginals.py @@ -26,6 +26,87 @@ def test_single_marginals(backend, px_fn, marginal, orientation): assert len(fig.data) == 1 + (marginal is not None) +@pytest.mark.parametrize("px_fn", [px.density_heatmap, px.density_contour]) +def test_marginal_heatmap_uses_z_and_histfunc(backend, px_fn): + 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_fn( + df, + x="total_bill", + y="tip", + z="size", + histfunc="sum", + marginal_x="heatmap", + marginal_y="heatmap", + ) + assert len(fig.data) == 3 + marginal_x_trace, marginal_y_trace = 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" + + +@pytest.mark.parametrize("px_fn", [px.density_heatmap, px.density_contour]) +def test_marginal_heatmap_without_z(backend, px_fn): + df = px.data.tips(return_type=backend) + + fig = px_fn( + df, x="total_bill", y="tip", marginal_x="heatmap", marginal_y="heatmap" + ) + marginal_x_trace, marginal_y_trace = 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" + + +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") + + +def test_marginal_heatmap_with_discrete_color(backend): # density_contour + line.color + df = px.data.tips(return_type=backend) + fig = px.density_contour( + df, x="total_bill", y="tip", color="sex", marginal_x="heatmap" + ) + assert len(fig.data) == 4 + assert [t.type for t in fig.data] == [ + "histogram2dcontour", + "histogram2d", + "histogram2dcontour", + "histogram2d", + ] + + def test_unsupported_marginal_raises_clear_error(): # issue 4654 # An unsupported marginal type used to fail deep inside make_figure with a # cryptic "'NoneType' object has no attribute 'constructor'". It should