multicomponent static covariate support for TFT

unit8co · dennisbader · Jun 5, 2022 · Apr 13, 2022 · Apr 23, 2022 · May 19, 2022
commit 3661385fe74bddff4c70a059e180fa99ceec9000
@@ -40,6 +40,7 @@ def __init__(
         self,
         output_dim: Tuple[int, int],
         variables_meta: Dict[str, Dict[str, List[str]]],
+        num_static_components: int,
         hidden_size: Union[int, List[int]] = 16,
         lstm_layers: int = 1,
         num_attention_heads: int = 4,
@@ -60,6 +61,9 @@ def __init__(
             shape of output given by (n_targets, loss_size). (loss_size corresponds to nr_params in other models).
         variables_meta : Dict[str, Dict[str, List[str]]]
             dict containing variable enocder, decoder variable names for mapping tensors in `_TFTModule.forward()`
+        num_static_components
+            the number of static components (not variables) of the input target series. This is either equal to the
+            number of target components or 1.
         hidden_size : int
             hidden state size of the TFT. It is the main hyper-parameter and common across the internal TFT
             architecture.
@@ -90,6 +94,7 @@ def __init__(
 
         self.n_targets, self.loss_size = output_dim
         self.variables_meta = variables_meta
+        self.num_static_components = num_static_components
         self.hidden_size = hidden_size
         self.hidden_continuous_size = hidden_continuous_size
         self.lstm_layers = lstm_layers
@@ -113,7 +118,11 @@ def __init__(
         # # processing inputs
         # continuous variable processing
         self.prescalers_linear = {
-            name: nn.Linear(1, self.hidden_continuous_size) for name in self.reals
+            name: nn.Linear(
+                1 if name not in self.static_variables else self.num_static_components,
+                self.hidden_continuous_size,
+            )
+            for name in self.reals
         }
 
         static_input_sizes = {
@@ -412,8 +421,7 @@ def forward(
         # Embedding and variable selection
         if self.static_variables:
             static_embedding = {
-                name: x_static[:, 0, i].unsqueeze(-1)
-                for i, name in enumerate(self.static_variables)
+                name: x_static[:, :, i] for i, name in enumerate(self.static_variables)
             }
             static_embedding, static_covariate_var = self.static_covariates_vsn(
                 static_embedding
@@ -864,9 +872,13 @@ def _create_model(self, train_sample: MixedCovariatesTrainTensorType) -> nn.Modu
             dict.fromkeys(static_input)
         )
 
+        n_static_components = (
+            len(static_covariates) if static_covariates is not None else 0
+        )
         return _TFTModule(
-            variables_meta=variables_meta,
             output_dim=self.output_dim,
+            variables_meta=variables_meta,
+            num_static_components=n_static_components,
             hidden_size=self.hidden_size,
             lstm_layers=self.lstm_layers,
             dropout=self.dropout,

@@ -2,7 +2,7 @@
 import pandas as pd
 import pytest
 
-from darts import TimeSeries
+from darts import TimeSeries, concatenate
 from darts.dataprocessing.transformers import Scaler
 from darts.logging import get_logger
 from darts.tests.base_test_class import DartsBaseTestClass
@@ -166,36 +166,39 @@ def test_mixed_covariates_and_accuracy(self):
             )
 
         def test_static_covariates_support(self):
-            target = tg.sine_timeseries(length=2, freq="h")
-            target = target.with_static_covariates(
-                pd.Series([0.0, 1.0], index=["st1", "st2"])
+            target_multi = concatenate(
+                [tg.sine_timeseries(length=10, freq="h")] * 2, axis=1
+            )
+
+            target_multi = target_multi.with_static_covariates(
+                pd.DataFrame([[0.0, 1.0], [2.0, 3.0]], index=["st1", "st2"])
             )
 
             # should work with cyclic encoding for time index
             model = TFTModel(
-                input_chunk_length=1,
-                output_chunk_length=1,
+                input_chunk_length=3,
+                output_chunk_length=4,
                 add_encoders={"cyclic": {"future": "hour"}},
                 pl_trainer_kwargs={"fast_dev_run": True},
             )
-            model.fit(target, verbose=False)
+            model.fit(target_multi, verbose=False)
             assert len(model.model.static_variables) == len(
-                target.static_covariates.columns
+                target_multi.static_covariates.columns
             )
 
-            model.predict(n=1, series=target, verbose=False)
+            model.predict(n=1, series=target_multi, verbose=False)
 
             # raise an error when trained with static covariates of wrong dimensionality
-            target = target.with_static_covariates(
-                pd.concat([target.static_covariates] * 2, axis=1)
+            target_multi = target_multi.with_static_covariates(
+                pd.concat([target_multi.static_covariates] * 2, axis=1)
             )
             with pytest.raises(ValueError):
-                model.predict(n=1, series=target, verbose=False)
+                model.predict(n=1, series=target_multi, verbose=False)
 
             # raise an error when trained with static covariates and trying to predict without
-            target = target.with_static_covariates(None)
+            target_multi = target_multi.with_static_covariates(None)
             with pytest.raises(ValueError):
-                model.predict(n=1, series=target, verbose=False)
+                model.predict(n=1, series=target_multi, verbose=False)
 
         def helper_generate_multivariate_case_data(self, season_length, n_repeat):
             """generates multivariate test case data. Target series is a sine wave stacked with a repeating