Merge pull request #48 from tvdboom/ts_plots

Ts plots

atom/atom.py

@@ -56,9 +56,9 @@
     FloatZeroToOneInc, Index, IndexSelector, Int, IntLargerEqualZero,
     IntLargerTwo, IntLargerZero, MetricConstructor, ModelsConstructor, NItems,
     NJobs, NormalizerStrats, NumericalStrats, Operators, Pandas, Predictor,
-    PrunerStrats, RowSelector, Scalar, ScalerStrats, Seasonality, Sequence,
-    Series, TargetSelector, Transformer, VectorizerStarts, Verbose, Warnings,
-    XSelector, YSelector, sequence_t,
+    PrunerStrats, RowSelector, Scalar, ScalerStrats, Seasonality,
+    SeasonalityMode, Sequence, Series, TargetSelector, Transformer,
+    VectorizerStarts, Verbose, Warnings, XSelector, YSelector, sequence_t,
 )
 from atom.utils.utils import (
     ClassMap, DataConfig, DataContainer, Goal, adjust_verbosity, bk,
@@ -922,7 +922,7 @@ def shrink(
             Whether to convert all features to sparse format. The value
             that is compressed is the most frequent value in the column.
 
-        columns: int, str, segment, sequence or None, default=None
+        columns: int, str, segment, sequence, dataframe or None, default=None
             [Selection of columns][row-and-column-selection] to shrink. If
             None, transform all columns.
 
@@ -1201,7 +1201,7 @@ def _add_transformer(
             has the `n_jobs` and/or `random_state` parameters, it
             adopts atom's values.
 
-        columns: int, str, segment, sequence or None, default=None
+        columns: int, str, segment, sequence, dataframe or None, default=None
             Columns in the dataset to transform. If None, transform
             all features.
 
@@ -1388,7 +1388,7 @@ def add(
             instance), and it has the `n_jobs` and/or `random_state`
             parameters, it adopts atom's values.
 
-        columns: int, str, segment, sequence or None, default=None
+        columns: int, str, segment, sequence, dataframe or None, default=None
             [Selection of columns][row-and-column-selection] to
             transform. Only select features or the target column, not
             both at the same time (if that happens, the target column
@@ -1564,7 +1564,7 @@ def decompose(
         self,
         *,
         model: str | Predictor | None = None,
-        mode: Literal["additive", "multiplicative"] = "additive",
+        mode: SeasonalityMode = "additive",
         **kwargs,
     ):
         """Detrend and deseasonalize the time series.
@@ -1584,9 +1584,7 @@ def decompose(
             * Use the `columns` parameter to only decompose the target
               column, e.g., `atom.decompose(columns=atom.target)`.
             * Use the [plot_decomposition][] method to visualize the
-              trend, seasonality and residuals of the time series. This
-              can help to determine if the data follows an additive or
-              multiplicative trend.
+              trend, seasonality and residuals of the time series.
 
         """
         columns = kwargs.pop("columns", None)

atom/data_cleaning.py

@@ -47,8 +47,9 @@
     Bins, Bool, CategoricalStrats, DataFrame, DiscretizerStrats, Engine,
     Estimator, FloatLargerZero, IntLargerEqualZero, IntLargerTwo,
     IntLargerZero, NJobs, NormalizerStrats, NumericalStrats, Pandas, Predictor,
-    PrunerStrats, Scalar, ScalerStrats, Sequence, Series, Transformer, Verbose,
-    XSelector, YSelector, dataframe_t, sequence_t, series_t,
+    PrunerStrats, Scalar, ScalerStrats, SeasonalityMode, Sequence, Series,
+    Transformer, Verbose, XSelector, YSelector, dataframe_t, sequence_t,
+    series_t,
 )
 from atom.utils.utils import (
     Goal, bk, composed, crash, get_col_order, get_cols, it, lst, merge,
@@ -1083,7 +1084,7 @@ def __init__(
         *,
         model: str | Predictor | None = None,
         sp: IntLargerZero | None = None,
-        mode: Literal["additive", "multiplicative"] = "additive",
+        mode: SeasonalityMode = "additive",
         n_jobs: NJobs = 1,
         verbose: Verbose = 0,
         logger: str | Path | Logger | None = None,

atom/plots/basefigure.py

@@ -190,7 +190,7 @@ def get_elem(
         else:
             return self.style[element].setdefault(name, next(getattr(self, element)))
 
-    def showlegend(self, name: str, legend: Legend | dict | None) -> bool:
+    def showlegend(self, name: str, legend: Legend | dict[str, Any] | None) -> bool:
         """Get whether the trace should be showed in the legend.
 
         If there's already a trace with the same name, it's not

atom/plots/baseplot.py

@@ -392,8 +392,8 @@ def _draw_line(
         child: str | None = None,
         legend: Legend | dict[str, Any] | None = None,
         **kwargs,
-    ) -> go.Scatter:
-        """Draw a line.
+    ):
+        """Draw a line on the current figure.
 
         Unify the style to draw a line, where parent and child
         (e.g., model - data set or column - distribution) keep the
@@ -408,19 +408,16 @@ def _draw_line(
         child: str or None, default=None
             Name of the secondary attribute.
 
-        legend: str, dict or None
+        legend: str, dict or None, default=None
             Legend argument provided by the user.
 
         **kwargs
             Additional keyword arguments for the trace.
 
-        Returns
-        -------
-        go.Scatter
-            New trace to add to figure.
-
         """
-        return go.Scatter(
+        BasePlot._fig.figure.add_scatter(
+            name=kwargs.pop("name", child or parent),
+            mode=kwargs.pop("mode", "lines"),
             line=kwargs.pop(
                 "line", {
                     "width": self.line_width,
@@ -440,15 +437,14 @@ def _draw_line(
                 "hovertemplate",
                 f"(%{{x}}, %{{y}})<extra>{parent}{f' - {child}' if child else ''}</extra>",
             ),
-            name=kwargs.pop("name", child or parent),
             legendgroup=kwargs.pop("legendgroup", parent),
             legendgrouptitle=kwargs.pop(
                 "legendgrouptitle",
                 {"text": parent, "font_size": self.label_fontsize} if child else None,
             ),
             showlegend=kwargs.pop(
                 "showlegend",
-                BasePlot._fig.showlegend(f"{parent}-{child}", legend)
+                BasePlot._fig.showlegend(f"{parent}-{child}" if child else parent, legend)
             ),
             **kwargs,
         )