fix lags_past_covariates dict type breaks lags_future_covariates when…

… output_chunk_shift>0 (unit8co#2655) * if clause modified to correct bug * contribution added to changelog * implemented modifications from reviewer * unit tests added --------- Co-authored-by: “authierj” <“[email protected]”>
ymatzkevich · Feb 6, 2025 · 06a80c6 · 06a80c6
1 parent c151dae
commit 06a80c6
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -59,6 +59,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Fixed**
 
+- Fixed a bug when initiating a `RegressionModel` with `lags_past_covariates` as dict and `lags_future_covariates` as some other type (not dict) and `output_chunk_shift>0`, [#2652](https://github.com/unit8co/darts/issues/2652) by [Jules Authier](https://github.com/authierj).
 - Fixed a bug which raised an error when computing residuals (or backtest with "per time step" metrics) on multiple series with corresponding historical forecasts of different lengths. [#2604](https://github.com/unit8co/darts/pull/2604) by [Dennis Bader](https://github.com/dennisbader).
 - Fixed a bug when using `darts.utils.data.tabularization.create_lagged_component_names()` with target `lags=None`, that did not return any lagged target label component names. [#2576](https://github.com/unit8co/darts/pull/2576) by [Dennis Bader](https://github.com/dennisbader).
 - Fixed a bug when using `num_samples > 1` with a deterministic regression model and the optimized `historical_forecasts()` method, which did not raise an exception. [#2576](https://github.com/unit8co/darts/pull/2588) by [Antoine Madrona](https://github.com/madtoinou).

diff --git a/darts/models/forecasting/regression_model.py b/darts/models/forecasting/regression_model.py
@@ -381,8 +381,11 @@ def _generate_lags(
                 else:
                     max_lags = max(max_lags, tmp_components_lags[comp_name][-1])
 
+            # Check if only default lags are provided
+            has_default_lags = list(tmp_components_lags.keys()) == ["default_lags"]
+
             # revert to shared lags logic when applicable
-            if list(tmp_components_lags.keys()) == ["default_lags"]:
+            if has_default_lags:
                 processed_lags[lags_abbrev] = tmp_components_lags["default_lags"]
             else:
                 processed_lags[lags_abbrev] = [min_lags, max_lags]
@@ -393,7 +396,7 @@ def _generate_lags(
                 processed_lags[lags_abbrev] = [
                     lag_ + output_chunk_shift for lag_ in processed_lags[lags_abbrev]
                 ]
-                if processed_component_lags:
+                if processed_component_lags and not has_default_lags:
                     processed_component_lags[lags_abbrev] = {
                         comp_: [lag_ + output_chunk_shift for lag_ in lags_]
                         for comp_, lags_ in processed_component_lags[

diff --git a/darts/tests/models/forecasting/test_regression_models.py b/darts/tests/models/forecasting/test_regression_models.py
@@ -2390,6 +2390,20 @@ def test_component_specific_lags(self, config):
                     "lags_past_covariates": 4,
                     "lags_future_covariates": [-3, 1],
                 },
+                # check that component-specific lags with output_chunk_shift works
+                {
+                    "lags_past_covariates": {"lin_past": [-3, -1]},
+                    "lags_future_covariates": [1, 2],
+                },
+                {
+                    "lags_past_covariates": [-3, -1],
+                    "lags_future_covariates": {"lin_future": [1, 2]},
+                },
+                {
+                    "lags": {"gaussian": 5},
+                    "lags_past_covariates": [-3, -1],
+                    "lags_future_covariates": [1, 2],
+                },
             ],
             [True, False],
             [3, 5],
@@ -2431,11 +2445,19 @@ def test_same_result_output_chunk_shift(self, config):
         )
         # adjusting the future lags should give identical models to non-shifted
         list_lags_adj = deepcopy(list_lags)
+        # this loop works for both component-specific and non-component-specific future lags
         if "lags_future_covariates" in list_lags_adj:
-            list_lags_adj["lags_future_covariates"] = [
-                lag_ - output_chunk_shift
-                for lag_ in list_lags_adj["lags_future_covariates"]
-            ]
+            if isinstance(list_lags_adj["lags_future_covariates"], dict):
+                for key in list_lags_adj["lags_future_covariates"]:
+                    list_lags_adj["lags_future_covariates"][key] = [
+                        lag_ - output_chunk_shift
+                        for lag_ in list_lags_adj["lags_future_covariates"][key]
+                    ]
+            else:
+                list_lags_adj["lags_future_covariates"] = [
+                    lag_ - output_chunk_shift
+                    for lag_ in list_lags_adj["lags_future_covariates"]
+                ]
         model_shift_adj = LinearRegressionModel(
             **list_lags_adj,
             output_chunk_shift=output_chunk_shift,