trendline_scope

Author: nicolaskruchten

packages/python/plotly/plotly/express/_chart_types.py

@@ -48,6 +48,7 @@ def scatter(
     trendline=None,
     trendline_options=None,
     trendline_color_override=None,
+    trendline_scope="trace",
     log_x=False,
     log_y=False,
     range_x=None,
@@ -93,6 +94,7 @@ def density_contour(
     trendline=None,
     trendline_options=None,
     trendline_color_override=None,
+    trendline_scope="trace",
     log_x=False,
     log_y=False,
     range_x=None,
@@ -202,7 +204,9 @@ def density_heatmap(
         z=[
             "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`.",],
+        histfunc=[
+            "The arguments to this function are the values of `z`.",
+        ],
     ),
 )
 
@@ -467,7 +471,9 @@ def histogram(
         args=locals(),
         constructor=go.Histogram,
         trace_patch=dict(
-            histnorm=histnorm, histfunc=histfunc, cumulative=dict(enabled=cumulative),
+            histnorm=histnorm,
+            histfunc=histfunc,
+            cumulative=dict(enabled=cumulative),
         ),
         layout_patch=dict(barmode=barmode, barnorm=barnorm),
     )
@@ -527,7 +533,11 @@ def violin(
         args=locals(),
         constructor=go.Violin,
         trace_patch=dict(
-            points=points, box=dict(visible=box), scalegroup=True, x0=" ", y0=" ",
+            points=points,
+            box=dict(visible=box),
+            scalegroup=True,
+            x0=" ",
+            y0=" ",
         ),
         layout_patch=dict(violinmode=violinmode),
     )

packages/python/plotly/plotly/express/_core.py

@@ -347,6 +347,9 @@ def make_trace_kwargs(args, trace_spec, trace_data, mapping_labels, sizeref):
                             )
 
                     # preserve original values of "x" in case they're dates
+                    # otherwise numpy/pandas can mess with the timezones
+                    # NB this means trendline functions must output one-to-one with the input series
+                    # i.e. we can't do resampling, because then the X values might not line up!
                     non_missing = np.logical_not(
                         np.logical_or(np.isnan(y), np.isnan(x))
                     )
@@ -867,23 +870,25 @@ def make_trace_spec(args, constructor, attrs, trace_patch):
             result.append(trace_spec)
 
     # Add trendline trace specifications
-    if "trendline" in args and args["trendline"]:
-        trace_spec = TraceSpec(
-            constructor=go.Scattergl
-            if constructor == go.Scattergl  # could be contour
-            else go.Scatter,
-            attrs=["trendline"],
-            trace_patch=dict(mode="lines"),
-            marginal=None,
-        )
-        if args["trendline_color_override"]:
-            trace_spec.trace_patch["line"] = dict(
-                color=args["trendline_color_override"]
-            )
-        result.append(trace_spec)
+    if args.get("trendline") and args.get("trendline_scope", "trace") == "trace":
+        result.append(make_trendline_spec(args, constructor))
     return result
 
 
+def make_trendline_spec(args, constructor):
+    trace_spec = TraceSpec(
+        constructor=go.Scattergl
+        if constructor == go.Scattergl  # could be contour
+        else go.Scatter,
+        attrs=["trendline"],
+        trace_patch=dict(mode="lines"),
+        marginal=None,
+    )
+    if args["trendline_color_override"]:
+        trace_spec.trace_patch["line"] = dict(color=args["trendline_color_override"])
+    return trace_spec
+
+
 def one_group(x):
     return ""
 
@@ -2127,6 +2132,27 @@ def make_figure(args, constructor, trace_patch=None, layout_patch=None):
         fig.update_layout(template=args["template"], overwrite=True)
     fig.frames = frame_list if len(frames) > 1 else []
 
+    if args.get("trendline") and args.get("trendline_scope", "trace") == "overall":
+        trendline_spec = make_trendline_spec(args, constructor)
+        trendline_trace = trendline_spec.constructor(
+            name="Overall Trendline", legendgroup="Overall Trendline", showlegend=False
+        )
+        if "line" not in trendline_spec.trace_patch:  # no color override
+            for m in grouped_mappings:
+                if m.variable == "color":
+                    next_color = m.sequence[len(m.val_map) % len(m.sequence)]
+                    trendline_spec.trace_patch["line"] = dict(color=next_color)
+        patch, fit_results = make_trace_kwargs(
+            args, trendline_spec, args["data_frame"], {}, sizeref
+        )
+        trendline_trace.update(patch)
+        fig.add_trace(
+            trendline_trace, row="all", col="all", exclude_empty_subplots=True
+        )
+        fig.update_traces(selector=-1, showlegend=True)
+        if fit_results is not None:
+            trendline_rows.append(dict(px_fit_results=fit_results))
+
     fig._px_trendlines = pd.DataFrame(trendline_rows)
 
     configure_axes(args, constructor, fig, orders)

packages/python/plotly/plotly/express/_doc.py

@@ -325,7 +325,10 @@
         "Setting this value is recommended when using `plotly.express.colors.diverging` color scales as the inputs to `color_continuous_scale`.",
     ],
     size_max=["int (default `20`)", "Set the maximum mark size when using `size`."],
-    markers=["boolean (default `False`)", "If `True`, markers are shown on lines.",],
+    markers=[
+        "boolean (default `False`)",
+        "If `True`, markers are shown on lines.",
+    ],
     log_x=[
         "boolean (default `False`)",
         "If `True`, the x-axis is log-scaled in cartesian coordinates.",
@@ -420,6 +423,10 @@
         "Valid CSS color.",
         "If provided, and if `trendline` is set, all trendlines will be drawn in this color rather than in the same color as the traces from which they draw their inputs.",
     ],
+    trendline_scope=[
+        "str (one of `'trace'` or `'overall'`, default `'trace'`)",
+        "If `'trace'`, then one trendline is drawn per trace (i.e. per color, symbol, facet, animation frame etc) and if `'overall'` then one trendline is computed for the entire dataset, and replicated across all facets.",
+    ],
     render_mode=[
         "str",
         "One of `'auto'`, `'svg'` or `'webgl'`, default `'auto'`",

packages/python/plotly/plotly/tests/test_optional/test_px/test_trendline.py

@@ -200,3 +200,42 @@ def test_trendline_on_timeseries(mode, options):
     assert type(fig.data[1].x[0]) == datetime
     assert np.all(fig.data[0].x == fig.data[1].x)
     assert str(fig.data[0].x[0]) == str(fig.data[1].x[0])
+
+
+def test_overall_trendline():
+    df = px.data.tips()
+    fig1 = px.scatter(df, x="total_bill", y="tip", trendline="ols")
+    assert len(fig1.data) == 2
+    assert "trendline" in fig1.data[1].hovertemplate
+    results1 = px.get_trendline_results(fig1)
+    params1 = results1["px_fit_results"].iloc[0].params
+
+    fig2 = px.scatter(
+        df,
+        x="total_bill",
+        y="tip",
+        color="sex",
+        trendline="ols",
+        trendline_scope="overall",
+    )
+    assert len(fig2.data) == 3
+    assert "trendline" in fig2.data[2].hovertemplate
+    results2 = px.get_trendline_results(fig2)
+    params2 = results2["px_fit_results"].iloc[0].params
+
+    assert np.all(np.array_equal(params1, params2))
+
+    fig3 = px.scatter(
+        df,
+        x="total_bill",
+        y="tip",
+        facet_row="sex",
+        trendline="ols",
+        trendline_scope="overall",
+    )
+    assert len(fig3.data) == 4
+    assert "trendline" in fig3.data[3].hovertemplate
+    results3 = px.get_trendline_results(fig3)
+    params3 = results3["px_fit_results"].iloc[0].params
+
+    assert np.all(np.array_equal(params1, params3))