Make test_ts optional in plot_forecast

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Add `ts.inverse_transform` as final step at `Pipeline.fit` method ([#316](https://github.com/tinkoff-ai/etna/pull/316))
+- Make test_ts optional in plot_forecast ([#321](https://github.com/tinkoff-ai/etna/pull/321))
 
 ## [1.3.3] - 2021-11-24
 ### Added

etna/analysis/plotters.py

@@ -20,7 +20,7 @@
 
 def plot_forecast(
     forecast_ts: "TSDataset",
-    test_ts: "TSDataset",
+    test_ts: Optional["TSDataset"] = None,
     train_ts: Optional["TSDataset"] = None,
     segments: Optional[List[str]] = None,
     n_train_samples: Optional[int] = None,
@@ -45,7 +45,7 @@ def plot_forecast(
         number of graphics columns
     """
     if not segments:
-        segments = list(set(test_ts.columns.get_level_values("segment")))
+        segments = list(set(forecast_ts.columns.get_level_values("segment")))
     segments_number = len(segments)
     columns_num = min(columns_num, len(segments))
     rows_num = math.ceil(segments_number / columns_num)
@@ -55,7 +55,8 @@ def plot_forecast(
 
     if train_ts is not None:
         train_ts.df.sort_values(by="timestamp", inplace=True)
-    test_ts.df.sort_values(by="timestamp", inplace=True)
+    if test_ts is not None:
+        test_ts.df.sort_values(by="timestamp", inplace=True)
     forecast_ts.df.sort_values(by="timestamp", inplace=True)
 
     for i, segment in enumerate(segments):
@@ -64,7 +65,10 @@ def plot_forecast(
         else:
             segment_train_df = pd.DataFrame(columns=["timestamp", "target", "segment"])
 
-        segment_test_df = test_ts[:, segment, :][segment]
+        if test_ts is not None:
+            segment_test_df = test_ts[:, segment, :][segment]
+        else:
+            segment_test_df = pd.DataFrame(columns=["timestamp", "target", "segment"])
 
         if n_train_samples is None:
             plot_df = segment_train_df
@@ -77,7 +81,8 @@ def plot_forecast(
 
         if (train_ts is not None) and (n_train_samples != 0):
             ax[i].plot(plot_df.index.values, plot_df.target.values, label="train")
-        ax[i].plot(segment_test_df.index.values, segment_test_df.target.values, label="test")
+        if test_ts is not None:
+            ax[i].plot(segment_test_df.index.values, segment_test_df.target.values, label="test")
         ax[i].plot(segment_forecast_df.index.values, segment_forecast_df.target.values, label="forecast")
         ax[i].set_title(segment)
         ax[i].tick_params("x", rotation=45)