lagged_regressors

ourownstory · Jan 4, 2024 · 660934c · 660934c
1 parent 381c912
commit 660934c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
@@ -124,7 +124,19 @@ def __getitem__(self, index):
             return sample, targets, meta
         else:
             start_idx = index
-            end_idx = start_idx + self.kwargs.get("n_lags") + self.kwargs.get("n_forecasts")
+
+            # Lagged Regressors
+            if self.kwargs["config_lagged_regressors"]:
+                n_lagged_regressor_list = []
+                for dict_name, nested_dict in self.kwargs["config_lagged_regressors"].items():
+                    name_of_nested_dict = dict_name
+                    n_lagged_regressor = self.kwargs["config_lagged_regressors"][name_of_nested_dict].n_lags
+                    n_lagged_regressor_list.append(n_lagged_regressor)
+                max_lag = max(self.kwargs["n_lags"], *n_lagged_regressor_list)
+                end_idx = start_idx + max_lag + self.kwargs.get("n_forecasts")
+
+            else:
+                end_idx = start_idx + self.kwargs.get("n_lags") + self.kwargs.get("n_forecasts")
 
             df_slice = self.df.iloc[start_idx:end_idx]
 

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -1713,7 +1713,6 @@ def test_on_the_fly_sampling():
     date_range = pd.date_range(start=start_date, end=end_date, freq="D")
     y = np.random.randint(0, 20, size=(len(date_range),))
     df = pd.DataFrame({"ds": date_range, "y": y})
-    df.loc[3, "y"] = np.nan
 
     m = NeuralProphet(epochs=1, learning_rate=0.01)
     m.fit(df, freq="H")