Extend STLDecomposer to Support Multiseries

docs/source/release_notes.rst

@@ -2,6 +2,7 @@ Release Notes
 -------------
 **Future Releases**
     * Enhancements
+        * Extended STLDecomposer to Support Multiseries :pr:`4253`
     * Fixes
     * Changes
     * Documentation Changes

evalml/pipelines/components/transformers/preprocessing/decomposer.py

@@ -324,7 +324,7 @@
     def plot_decomposition(
         self,
         X: pd.DataFrame,
-        y: pd.Series,
+        y: pd.DataFrame,
         show: bool = False,
     ) -> tuple[plt.Figure, list]:
         """Plots the decomposition of the target signal.
@@ -340,20 +340,46 @@
                 plotted on them
 
         """
+        if isinstance(y, pd.Series):
+            y = y.to_frame()
+
+        plot_info = {}
+        if self.frequency and self.time_index and len(y.columns) > 1:
+            X.index = pd.DatetimeIndex(X[self.time_index], freq=self.frequency)
         decomposition_results = self.get_trend_dataframe(X, y)
-        fig, axs = plt.subplots(4)
-        fig.set_size_inches(18.5, 14.5)
-        axs[0].plot(decomposition_results[0]["signal"], "r")
-        axs[0].set_title("signal")
-        axs[1].plot(decomposition_results[0]["trend"], "b")
-        axs[1].set_title("trend")
-        axs[2].plot(decomposition_results[0]["seasonality"], "g")
-        axs[2].set_title("seasonality")
-        axs[3].plot(decomposition_results[0]["residual"], "y")
-        axs[3].set_title("residual")
-        if show:  # pragma: no cover
-            plt.show()
-        return fig, axs
+
+        # Iterate through each series id
+        for id in y.columns:
+            fig, axs = plt.subplots(4)
+            fig.set_size_inches(18.5, 14.5)
+
+            for ax in axs:
+                ax.cla()
+
+            if len(y.columns) > 1:
+                results = decomposition_results[id]
+            else:
+                results = decomposition_results
+            axs[0].plot(results[0]["signal"], "r")
+            axs[0].set_title("signal")
+            axs[1].plot(results[0]["trend"], "b")
+            axs[1].set_title("trend")
+            axs[2].plot(results[0]["seasonality"], "g")
+            axs[2].set_title("seasonality")
+            axs[3].plot(results[0]["residual"], "y")
+            axs[3].set_title("residual")
+
+            # If multiseries, return a dictionary of tuples
+            if len(y.columns) > 1:
+                fig.suptitle("Decomposition for Series {}".format(id))
+                plot_info[id] = (fig, axs)
+            else:
+                plot_info = (fig, axs)
+
+            if show:  # pragma: no cover
+                plt.show()
+
+        return plot_info
 
     def _check_target(self, X: pd.DataFrame, y: pd.Series):
         """Function to ensure target is not None and has a pandas.DatetimeIndex."""