Fix/dlinear and nlinear use_static_cov. with multivariate series (#2070)

* feat: added test to check that use_static_covariates covers all possible static covariates representations

* fix: properly account for the two possible static covariates representation in multivariates series

* fix: typo in the warning message

* feat: added type hint, reordered docstring to match argument order

* feat: added type hint, reordered docstring to match argument order (dlinear)

* feat: updated changelog

---------

Co-authored-by: Dennis Bader <[email protected]>

Author: madtoinou

CHANGELOG.md

@@ -31,6 +31,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 - Fixed a bug when loading a `TorchForecastingModel` that was trained with a precision other than `float64`. [#2046](https://github.com/unit8co/darts/pull/2046) by [Freddie Hsin-Fu Huang](https://github.com/Hsinfu).
 - Fixed broken links in the `Transfer learning` example notebook with publicly hosted version of the three datasets. [#2067](https://github.com/unit8co/darts/pull/2067) by [Antoine Madrona](https://github.com/madtoinou).
 - Fixed a bug when using `NLinearModel` on multivariate series with covariates and `normalize=True`. [#2072](https://github.com/unit8co/darts/pull/2072) by [Antoine Madrona](https://github.com/madtoinou).
+- Fixed a bug when using `DLinearModel` and `NLinearModel` on multivariate series with "components-shared" static covariates and `use_static_covariates=True`. [#2070](https://github.com/unit8co/darts/pull/2070) by [Antoine Madrona](https://github.com/madtoinou).
 
 ### For developers of the library:
 

darts/models/forecasting/dlinear.py

@@ -68,14 +68,14 @@ class _DLinearModule(PLMixedCovariatesModule):
 
     def __init__(
         self,
-        input_dim,
-        output_dim,
-        future_cov_dim,
-        static_cov_dim,
-        nr_params,
-        shared_weights,
-        kernel_size,
-        const_init,
+        input_dim: int,
+        output_dim: int,
+        future_cov_dim: int,
+        static_cov_dim: int,
+        nr_params: int,
+        shared_weights: bool,
+        kernel_size: int,
+        const_init: bool,
         **kwargs,
     ):
         """PyTorch module implementing the DLinear architecture.
@@ -89,7 +89,7 @@ def __init__(
         future_cov_dim
             Number of components in the future covariates
         static_cov_dim
-            Dimensionality of the static covariates
+            Dimensionality of the static covariates (either component-specific or shared)
         nr_params
             The number of parameters of the likelihood (or 1 if no likelihood is used).
         shared_weights
@@ -113,8 +113,6 @@ def __init__(
             Tensor containing the output of the NBEATS module.
         """
 
-        # TODO: could we support future covariates with a simple extension?
-
         super().__init__(**kwargs)
         self.input_dim = input_dim
         self.output_dim = output_dim
@@ -142,9 +140,6 @@ def _create_linear_layer(in_dim, out_dim):
             layer_in_dim = self.input_chunk_length * self.input_dim
             layer_out_dim = self.output_chunk_length * self.output_dim * self.nr_params
 
-            # for static cov, we take the number of components of the target, times static cov dim
-            layer_in_dim_static_cov = self.output_dim * self.static_cov_dim
-
         self.linear_seasonal = _create_linear_layer(layer_in_dim, layer_out_dim)
         self.linear_trend = _create_linear_layer(layer_in_dim, layer_out_dim)
 
@@ -155,7 +150,7 @@ def _create_linear_layer(in_dim, out_dim):
             )
         if self.static_cov_dim != 0:
             self.linear_static_cov = _create_linear_layer(
-                layer_in_dim_static_cov, layer_out_dim
+                self.static_cov_dim, layer_out_dim
             )
 
     @io_processor
@@ -477,8 +472,8 @@ def _create_model(
         raise_if(
             self.shared_weights
             and (train_sample[1] is not None or train_sample[2] is not None),
-            "Covariates have been provided, but the model has been built with shared_weights=True."
-            + "Please set shared_weights=False to use covariates.",
+            "Covariates have been provided, but the model has been built with shared_weights=True. "
+            "Please set shared_weights=False to use covariates.",
         )
 
         input_dim = train_sample[0].shape[1] + sum(
@@ -488,8 +483,11 @@ def _create_model(
         )
         future_cov_dim = train_sample[3].shape[1] if train_sample[3] is not None else 0
 
-        # dimension is (component, static_dim), we extract static_dim
-        static_cov_dim = train_sample[4].shape[1] if train_sample[4] is not None else 0
+        if train_sample[4] is None:
+            static_cov_dim = 0
+        else:
+            # account for component-specific or shared static covariates representation
+            static_cov_dim = train_sample[4].shape[0] * train_sample[4].shape[1]
 
         output_dim = train_sample[-1].shape[1]
         nr_params = 1 if self.likelihood is None else self.likelihood.num_parameters

darts/models/forecasting/nlinear.py

@@ -23,14 +23,14 @@ class _NLinearModule(PLMixedCovariatesModule):
 
     def __init__(
         self,
-        input_dim,
-        output_dim,
-        future_cov_dim,
-        static_cov_dim,
-        nr_params,
-        shared_weights,
-        const_init,
-        normalize,
+        input_dim: int,
+        output_dim: int,
+        future_cov_dim: int,
+        static_cov_dim: int,
+        nr_params: int,
+        shared_weights: bool,
+        const_init: bool,
+        normalize: bool,
         **kwargs,
     ):
         """PyTorch module implementing the N-HiTS architecture.
@@ -44,16 +44,16 @@ def __init__(
         future_cov_dim
             Number of components in the future covariates
         static_cov_dim
-            Dimensionality of the static covariates
+            Dimensionality of the static covariates (either component-specific or shared)
         nr_params
             The number of parameters of the likelihood (or 1 if no likelihood is used).
         shared_weights
             Whether to use shared weights for the components of the series.
             ** Ignores covariates when True. **
-        normalize
-            Whether to apply the "normalization" described in the paper.
         const_init
             Whether to initialize the weights to 1/in_len
+        normalize
+            Whether to apply the "normalization" described in the paper.
 
         **kwargs
             all parameters required for :class:`darts.model.forecasting_models.PLForecastingModule` base class.
@@ -94,9 +94,6 @@ def _create_linear_layer(in_dim, out_dim):
             layer_in_dim = self.input_chunk_length * self.input_dim
             layer_out_dim = self.output_chunk_length * self.output_dim * self.nr_params
 
-            # for static cov, we take the number of components of the target, times static cov dim
-            layer_in_dim_static_cov = self.output_dim * self.static_cov_dim
-
         self.layer = _create_linear_layer(layer_in_dim, layer_out_dim)
 
         if self.future_cov_dim != 0:
@@ -106,7 +103,7 @@ def _create_linear_layer(in_dim, out_dim):
             )
         if self.static_cov_dim != 0:
             self.linear_static_cov = _create_linear_layer(
-                layer_in_dim_static_cov, layer_out_dim
+                self.static_cov_dim, layer_out_dim
             )
 
     @io_processor
@@ -438,8 +435,11 @@ def _create_model(self, train_sample: Tuple[torch.Tensor]) -> torch.nn.Module:
         )
         future_cov_dim = train_sample[3].shape[1] if train_sample[3] is not None else 0
 
-        # dimension is (component, static_dim), we extract static_dim
-        static_cov_dim = train_sample[4].shape[1] if train_sample[4] is not None else 0
+        if train_sample[4] is None:
+            static_cov_dim = 0
+        else:
+            # account for component-specific or shared static covariates representation
+            static_cov_dim = train_sample[4].shape[0] * train_sample[4].shape[1]
 
         output_dim = train_sample[-1].shape[1]
 

darts/tests/models/forecasting/test_global_forecasting_models.py

@@ -1,5 +1,6 @@
 import os
 from copy import deepcopy
+from itertools import product
 from unittest.mock import ANY, patch
 
 import numpy as np
@@ -206,6 +207,15 @@ class TestGlobalForecastingModels:
         target = sine_1_ts + sine_2_ts + linear_ts + sine_3_ts
         target_past, target_future = target.split_after(split_ratio)
 
+        # various ts with different static covariates representations
+        ts_w_static_cov = tg.linear_timeseries(length=80).with_static_covariates(
+            pd.Series([1, 2])
+        )
+        ts_shared_static_cov = ts_w_static_cov.stack(tg.sine_timeseries(length=80))
+        ts_comps_static_cov = ts_shared_static_cov.with_static_covariates(
+            pd.DataFrame([[0, 1], [2, 3]], columns=["st1", "st2"])
+        )
+
         @pytest.mark.parametrize("config", models_cls_kwargs_errs)
         def test_save_model_parameters(self, config):
             # model creation parameters were saved before. check if re-created model has same params as original
@@ -450,6 +460,37 @@ def test_future_covariates(self):
             with pytest.raises(ValueError):
                 model.predict(n=161, future_covariates=self.covariates)
 
+        @pytest.mark.parametrize(
+            "model_cls,ts",
+            product(
+                [TFTModel, DLinearModel, NLinearModel, TiDEModel],
+                [ts_w_static_cov, ts_shared_static_cov, ts_comps_static_cov],
+            ),
+        )
+        def test_use_static_covariates(self, model_cls, ts):
+            """
+            Check that both static covariates representations are supported (component-specific and shared)
+            for both uni- and multivariate series when fitting the model.
+            Also check that the static covariates are present in the forecasted series
+            """
+            model = model_cls(
+                input_chunk_length=IN_LEN,
+                output_chunk_length=OUT_LEN,
+                random_state=0,
+                use_static_covariates=True,
+                n_epochs=1,
+                **tfm_kwargs,
+            )
+            # must provide mandatory future_covariates to TFTModel
+            model.fit(
+                series=ts,
+                future_covariates=self.sine_1_ts
+                if model.supports_future_covariates
+                else None,
+            )
+            pred = model.predict(OUT_LEN)
+            assert pred.static_covariates.equals(ts.static_covariates)
+
         def test_batch_predictions(self):
             # predicting multiple time series at once needs to work for arbitrary batch sizes
             # univariate case

darts/tests/models/forecasting/test_torch_forecasting_model.py

@@ -66,7 +66,7 @@
 
     TORCH_AVAILABLE = True
 except ImportError:
-    logger.warning("Torch not available. RNN tests will be skipped.")
+    logger.warning("Torch not available. Tests will be skipped.")
     TORCH_AVAILABLE = False
 
 if TORCH_AVAILABLE: