move trendline code to own module

Author: nicolaskruchten

doc/apidoc/plotly.express.rst

@@ -59,3 +59,4 @@ plotly's high-level API for rapid figure generation. ::
 
    generated/plotly.express.data.rst
    generated/plotly.express.colors.rst
+   generated/plotly.express.trendline_functions.rst

packages/python/plotly/plotly/express/__init__.py

@@ -59,7 +59,7 @@
 
 from ._special_inputs import IdentityMap, Constant, Range  # noqa: F401
 
-from . import data, colors  # noqa: F401
+from . import data, colors, trendline_functions  # noqa: F401
 
 __all__ = [
     "scatter",

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

@@ -2,6 +2,7 @@
 import plotly.io as pio
 from collections import namedtuple, OrderedDict
 from ._special_inputs import IdentityMap, Constant, Range
+from .trendline_functions import ols, lowess, ma, ewm
 
 from _plotly_utils.basevalidators import ColorscaleValidator
 from plotly.colors import qualitative, sequential
@@ -229,65 +230,6 @@ def make_mapping(args, variable):
     )
 
 
-def lowess(options, x, y, x_label, y_label, non_missing):
-    import statsmodels.api as sm
-
-    frac = options.get("frac", 0.6666666)
-    # missing ='drop' is the default value for lowess but not for OLS (None)
-    # we force it here in case statsmodels change their defaults
-    y_out = sm.nonparametric.lowess(y, x, missing="drop", frac=frac)[:, 1]
-    hover_header = "<b>LOWESS trendline</b><br><br>"
-    return y_out, hover_header, None
-
-
-def ma(options, x, y, x_label, y_label, non_missing):
-    y_out = pd.Series(y, index=x).rolling(**options).mean()[non_missing]
-    hover_header = "<b>Moving Average trendline</b><br><br>"
-    return y_out, hover_header, None
-
-
-def ewm(options, x, y, x_label, y_label, non_missing):
-    y_out = pd.Series(y, index=x).ewm(**options).mean()[non_missing]
-    hover_header = "<b>EWM trendline</b><br><br>"
-    return y_out, hover_header, None
-
-
-def ols(options, x, y, x_label, y_label, non_missing):
-    import statsmodels.api as sm
-
-    add_constant = options.get("add_constant", True)
-    log_x = options.get("log_x", False)
-    log_y = options.get("log_y", False)
-
-    if log_y:
-        y = np.log(y)
-    if log_x:
-        x = np.log(x)
-    if add_constant:
-        x = sm.add_constant(x)
-    fit_results = sm.OLS(y, x, missing="drop").fit()
-    y_out = fit_results.predict()
-    if log_y:
-        y_out = np.exp(y_out)
-    hover_header = "<b>OLS trendline</b><br>"
-    if len(fit_results.params) == 2:
-        hover_header += "%s = %g * %s + %g<br>" % (
-            y_label,
-            fit_results.params[1],
-            x_label,
-            fit_results.params[0],
-        )
-    elif not add_constant:
-        hover_header += "%s = %g* %s<br>" % (y_label, fit_results.params[0], x_label,)
-    else:
-        hover_header += "%s = %g<br>" % (y_label, fit_results.params[0],)
-    hover_header += "R<sup>2</sup>=%f<br><br>" % fit_results.rsquared
-    return y_out, hover_header, fit_results
-
-
-trendline_functions = dict(lowess=lowess, ma=ma, ewm=ewm, ols=ols)
-
-
 def make_trace_kwargs(args, trace_spec, trace_data, mapping_labels, sizeref):
     """Populates a dict with arguments to update trace
 
@@ -361,6 +303,7 @@ def make_trace_kwargs(args, trace_spec, trace_data, mapping_labels, sizeref):
                 if trace_spec.constructor == go.Histogram:
                     mapping_labels["count"] = "%{x}"
             elif attr_name == "trendline":
+                trendline_functions = dict(lowess=lowess, ma=ma, ewm=ewm, ols=ols)
                 if (
                     attr_value in trendline_functions
                     and args["x"]

packages/python/plotly/plotly/express/trendline_functions.py

@@ -0,0 +1,58 @@
+import pandas as pd
+import numpy as np
+
+
+def ols(options, x, y, x_label, y_label, non_missing):
+    import statsmodels.api as sm
+
+    add_constant = options.get("add_constant", True)
+    log_x = options.get("log_x", False)
+    log_y = options.get("log_y", False)
+
+    if log_y:
+        y = np.log(y)
+        y_label = "log(%s)" % y_label
+    if log_x:
+        x = np.log(x)
+        x_label = "log(%s)" % x_label
+    if add_constant:
+        x = sm.add_constant(x)
+    fit_results = sm.OLS(y, x, missing="drop").fit()
+    y_out = fit_results.predict()
+    if log_y:
+        y_out = np.exp(y_out)
+    hover_header = "<b>OLS trendline</b><br>"
+    if len(fit_results.params) == 2:
+        hover_header += "%s = %g * %s + %g<br>" % (
+            y_label,
+            fit_results.params[1],
+            x_label,
+            fit_results.params[0],
+        )
+    elif not add_constant:
+        hover_header += "%s = %g* %s<br>" % (y_label, fit_results.params[0], x_label,)
+    else:
+        hover_header += "%s = %g<br>" % (y_label, fit_results.params[0],)
+    hover_header += "R<sup>2</sup>=%f<br><br>" % fit_results.rsquared
+    return y_out, hover_header, fit_results
+
+
+def lowess(options, x, y, x_label, y_label, non_missing):
+    import statsmodels.api as sm
+
+    frac = options.get("frac", 0.6666666)
+    y_out = sm.nonparametric.lowess(y, x, missing="drop", frac=frac)[:, 1]
+    hover_header = "<b>LOWESS trendline</b><br><br>"
+    return y_out, hover_header, None
+
+
+def ma(options, x, y, x_label, y_label, non_missing):
+    y_out = pd.Series(y, index=x).rolling(**options).mean()[non_missing]
+    hover_header = "<b>Moving Average trendline</b><br><br>"
+    return y_out, hover_header, None
+
+
+def ewm(options, x, y, x_label, y_label, non_missing):
+    y_out = pd.Series(y, index=x).ewm(**options).mean()[non_missing]
+    hover_header = "<b>EWM trendline</b><br><br>"
+    return y_out, hover_header, None