dataset adaption for static covs

unit8co · dennisbader · Jun 5, 2022 · Apr 13, 2022 · Apr 23, 2022 · May 19, 2022
commit fde974e68908282ecdea0caa87ce14107a944306
@@ -31,7 +31,7 @@
 logger = get_logger(__name__)
 
 MixedCovariatesTrainTensorType = Tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
 ]
 
 
@@ -415,7 +415,7 @@ def forward(self, x: Tuple[torch.Tensor, Optional[torch.Tensor]]) -> torch.Tenso
 
         # Embedding and variable selection
         if static_covariates is not None:
-            # TODO: impelement static covariates
+            # TODO: implement static covariates
             # # static embeddings will be constant over entire batch
             # static_embedding = {name: input_vectors[name][:, 0] for name in self.static_variables}
             # static_embedding, static_covariate_var = self.static_covariates_vsn(static_embedding)
@@ -775,6 +775,7 @@ def _create_model(self, train_sample: MixedCovariatesTrainTensorType) -> nn.Modu
             historic_future_covariate,
             future_covariate,
             future_target,
+            static_covariates,
         ) = train_sample
 
         # add a covariate placeholder so that relative index will be included
@@ -800,7 +801,7 @@ def _create_model(self, train_sample: MixedCovariatesTrainTensorType) -> nn.Modu
                 axis=1,
             )
 
-        static_covariates = None  # placeholder for future
+        # static_covariates = None  # placeholder for future
 
         self.output_dim = (
             (future_target.shape[1], 1)

@@ -846,7 +846,7 @@ def fit_from_dataset(
         train_sample = train_dataset[0]
         if self.model is None:
             # Build model, based on the dimensions of the first series in the train set.
-            self.train_sample, self.output_dim = train_sample, train_sample[-1].shape[1]
+            self.train_sample, self.output_dim = train_sample, train_sample[-2].shape[1]
             self._init_model(trainer)
         else:
             # Check existing model has input/output dims matching what's provided in the training set.

@@ -233,10 +233,22 @@ def __len__(self):
 
     def __getitem__(
         self, idx
-    ) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray], np.ndarray]:
-        past_target, past_covariate, future_target = self.ds_past[idx]
-        _, future_covariate, _ = self.ds_future[idx]
-        return past_target, past_covariate, future_covariate, future_target
+    ) -> Tuple[
+        np.ndarray,
+        Optional[np.ndarray],
+        Optional[np.ndarray],
+        np.ndarray,
+        Optional[np.ndarray],
+    ]:
+        past_target, past_covariate, future_target, static_covariate = self.ds_past[idx]
+        _, future_covariate, _, static_covariate = self.ds_future[idx]
+        return (
+            past_target,
+            past_covariate,
+            future_covariate,
+            future_target,
+            static_covariate,
+        )
 
 
 class MixedCovariatesSequentialDataset(MixedCovariatesTrainingDataset):
@@ -322,16 +334,24 @@ def __getitem__(
         Optional[np.ndarray],
         Optional[np.ndarray],
         np.ndarray,
+        Optional[np.ndarray],
     ]:
 
-        past_target, past_covariate, future_target = self.ds_past[idx]
-        _, historic_future_covariate, future_covariate, _ = self.ds_dual[idx]
+        past_target, past_covariate, future_target, static_covariate = self.ds_past[idx]
+        (
+            _,
+            historic_future_covariate,
+            future_covariate,
+            _,
+            static_covariate,
+        ) = self.ds_dual[idx]
         return (
             past_target,
             past_covariate,
             historic_future_covariate,
             future_covariate,
             future_target,
+            static_covariate,
         )
 
 

@@ -518,7 +518,9 @@ def __init__(
     def __len__(self):
         return self.ideal_nr_samples
 
-    def __getitem__(self, idx) -> Tuple[np.ndarray, Optional[np.ndarray], np.ndarray]:
+    def __getitem__(
+        self, idx
+    ) -> Tuple[np.ndarray, Optional[np.ndarray], np.ndarray, Optional[np.ndarray]]:
         # determine the index of the time series.
         ts_idx = idx // self.max_samples_per_ts
         ts_target = self.target_series[ts_idx]
@@ -598,4 +600,9 @@ def __getitem__(self, idx) -> Tuple[np.ndarray, Optional[np.ndarray], np.ndarray
                 f"target series.",
             )
 
-        return past_target, covariate, future_target
+        static_covariate = (
+            ts_target.static_covariates.values
+            if ts_target.static_covariates is not None
+            else None
+        )
+        return past_target, covariate, future_target, static_covariate