Added plot param of the NeuralProphet fit() and uncomment test_progre…

…ss_display in test_integration.py. (#958) * Added plot param of the NeuralProphet fit(). * Applied black formatting to forecaster.py. Co-authored-by: Kevin R. Chen <[email protected]>
ourownstory · Nov 15, 2022 · 686335a · 686335a
1 parent 974ee77
commit 686335a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
@@ -635,7 +635,9 @@ def add_seasonality(self, name, period, fourier_order):
         self.config_season.append(name=name, period=period, resolution=fourier_order, arg="custom")
         return self
 
-    def fit(self, df, freq="auto", validation_df=None, progress="bar", minimal=False, continue_training=False):
+    def fit(
+        self, df, freq="auto", validation_df=None, progress="bar", minimal=False, continue_training=False, plot=True
+    ):
         """Train, and potentially evaluate model.
 
         Training/validation metrics may be distorted in case of auto-regression,
@@ -657,6 +659,8 @@ def fit(self, df, freq="auto", validation_df=None, progress="bar", minimal=False
                 whether to train without any printouts or metrics collection
             continue_training : bool
                 whether to continue training from the last checkpoint
+            plot : bool
+                where to show the progress plot or not
 
         Returns
         -------
@@ -725,7 +729,8 @@ def fit(self, df, freq="auto", validation_df=None, progress="bar", minimal=False
                 _ = plt.plot(metrics_df[["Loss"]])
             else:
                 _ = plt.plot(metrics_df[["Loss", "Loss_val"]])
-            plt.show()
+            if plot:
+                plt.show()
 
         self.fitted = True
         return metrics_df

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -1433,7 +1433,6 @@ def test_metrics():
     forecast = m2.predict(df)
 
 
-"""
 def test_progress_display():
     log.info("testing: Progress Display")
     df = pd.read_csv(AIR_FILE, nrows=100)
@@ -1445,8 +1444,7 @@ def test_progress_display():
             batch_size=BATCH_SIZE,
             learning_rate=LR,
         )
-        metrics_df = m.fit(df, progress=progress)
-"""
+        metrics_df = m.fit(df, progress=progress, plot=PLOT)
 
 
 def test_n_lags_for_regressors():