Fix QuantileForecast.plot() to use DateTimeIndex (#2269)

awslabs · Sep 5, 2022 · 2c946c6 · 2c946c6
1 parent 44ef129
commit 2c946c6
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt
@@ -7,3 +7,4 @@ ujson
 orjson
 requests
 holidays~=0.9
+matplotlib
diff --git a/src/gluonts/model/forecast.py b/src/gluonts/model/forecast.py
@@ -779,14 +779,11 @@ def plot(self, label=None, output_file=None, keys=None, *args, **kwargs):
             keys = self.forecast_keys
 
         for k, v in zip(keys, self.forecast_array):
-            plt.plot(
-                self.index,
-                v,
+            pd.Series(data=v, index=self.index.to_timestamp()).plot(
                 label=f"{label_prefix}q{k}",
                 *args,
                 **kwargs,
             )
-            plt.legend()
         if output_file:
             plt.savefig(output_file)
 

diff --git a/test/model/test_forecast.py b/test/model/test_forecast.py
@@ -60,6 +60,8 @@ def percentile(value):
     assert len(forecast.index) == pred_length
     assert forecast.index[0] == START_DATE
 
+    forecast.plot()
+
 
 @pytest.mark.parametrize(
     "forecast, exp_index",