[ENH] Tide model in v2 interface (#1889)

Adds `Tide` model from `dsipts` to `ptf-v2`

Created a new folder `tide_dsipts` in `tide` that contains all the
necessary parts for the tide

Author: phoeenniixx

pytorch_forecasting/data/data_module.py

@@ -430,8 +430,8 @@ def __getitem__(self, idx):
             encoder_indices = slice(start_idx, start_idx + enc_length)
             decoder_indices = slice(start_idx + enc_length, end_idx)
 
-            target_scale = data["target"][encoder_indices]
-            target_scale = target_scale[~torch.isnan(target_scale)].abs().mean()
+            target_past = data["target"][encoder_indices]
+            target_scale = target_past[~torch.isnan(target_past)].abs().mean()
             if torch.isnan(target_scale) or target_scale == 0:
                 target_scale = torch.tensor(1.0)
 
@@ -503,6 +503,7 @@ def __getitem__(self, idx):
                 "decoder_lengths": torch.tensor(pred_length),
                 "decoder_target_lengths": torch.tensor(pred_length),
                 "groups": data["group"],
+                "target_past": target_past,
                 "encoder_time_idx": torch.arange(enc_length),
                 "decoder_time_idx": torch.arange(enc_length, enc_length + pred_length),
                 "target_scale": target_scale,
@@ -713,6 +714,7 @@ def collate_fn(batch):
                 [x["decoder_target_lengths"] for x, _ in batch]
             ),
             "groups": torch.stack([x["groups"] for x, _ in batch]),
+            "target_past": torch.stack([x["target_past"] for x, _ in batch]),
             "encoder_time_idx": torch.stack([x["encoder_time_idx"] for x, _ in batch]),
             "decoder_time_idx": torch.stack([x["decoder_time_idx"] for x, _ in batch]),
             "target_scale": torch.stack([x["target_scale"] for x, _ in batch]),

pytorch_forecasting/layers/__init__.py

@@ -7,11 +7,13 @@
     FullAttention,
     TriangularCausalMask,
 )
+from pytorch_forecasting.layers._blocks import ResidualBlock
 from pytorch_forecasting.layers._decomposition import SeriesDecomposition
 from pytorch_forecasting.layers._embeddings import (
     DataEmbedding_inverted,
     EnEmbedding,
     PositionalEmbedding,
+    embedding_cat_variables,
 )
 from pytorch_forecasting.layers._encoders import (
     Encoder,
@@ -48,4 +50,6 @@
     "sLSTMLayer",
     "sLSTMNetwork",
     "SeriesDecomposition",
+    "ResidualBlock",
+    "embedding_cat_variables",
 ]

pytorch_forecasting/layers/_blocks/__init__.py

@@ -0,0 +1,3 @@
+from pytorch_forecasting.layers._blocks._residual_block_dsipts import ResidualBlock
+
+__all__ = ["ResidualBlock"]

pytorch_forecasting/layers/_blocks/_residual_block_dsipts.py

@@ -0,0 +1,50 @@
+import torch.nn as nn
+
+
+class ResidualBlock(nn.Module):
+    def __init__(
+        self, in_size: int, out_size: int, dropout_rate: float, activation_fun: str = ""
+    ):
+        """Residual Block as basic layer of the archetecture.
+
+        MLP with one hidden layer, activation and skip connection
+        Basically dimension d_model, but better if input_dim and output_dim are explicit
+
+        in_size and out_size to handle dimensions at different stages of the NN
+
+        Parameters
+        ----------
+        in_size: int
+            input size
+        out_size: int
+            output size
+        dropout_rate: float
+            dropout
+        activation_fun: str, Optional
+            activation function to use in the Residual Block. Defaults to nn.ReLU.
+        """  # noqa: E501
+        import ast
+
+        super().__init__()
+
+        self.direct_linear = nn.Linear(in_size, out_size, bias=False)
+
+        if activation_fun == "":
+            self.act = nn.ReLU()
+        else:
+            activation = ast.literal_eval(activation_fun)
+            self.act = activation()
+        self.lin = nn.Linear(in_size, out_size)
+        self.dropout = nn.Dropout(dropout_rate)
+
+        self.final_norm = nn.LayerNorm(out_size)
+
+    def forward(self, x, apply_final_norm=True):
+        direct_x = self.direct_linear(x)
+
+        x = self.dropout(self.lin(self.act(x)))
+
+        out = x + direct_x
+        if apply_final_norm:
+            return self.final_norm(out)
+        return out

pytorch_forecasting/layers/_embeddings/__init__.py

@@ -9,5 +9,11 @@
 from pytorch_forecasting.layers._embeddings._positional_embedding import (
     PositionalEmbedding,
 )
+from pytorch_forecasting.layers._embeddings._sub_nn import embedding_cat_variables
 
-__all__ = ["PositionalEmbedding", "DataEmbedding_inverted", "EnEmbedding"]
+__all__ = [
+    "PositionalEmbedding",
+    "DataEmbedding_inverted",
+    "EnEmbedding",
+    "embedding_cat_variables",
+]

pytorch_forecasting/layers/_embeddings/_sub_nn.py

@@ -0,0 +1,101 @@
+from typing import Union
+
+import torch
+import torch.nn as nn
+
+
+class embedding_cat_variables(nn.Module):
+    # at the moment cat_past and cat_fut together
+    def __init__(self, seq_len: int, lag: int, d_model: int, emb_dims: list, device):
+        """Class for embedding categorical variables, adding 3 positional variables during forward
+
+        Parameters
+        ----------
+        seq_len: int
+            length of the sequence (sum of past and future steps)
+        lag: (int):
+            number of future step to be predicted
+        hiden_size: int
+            dimension of all variables after they are embedded
+        emb_dims: list
+            size of the dictionary for embedding. One dimension for each categorical variable
+        device : torch.device
+        """  # noqa: E501
+        super().__init__()
+        self.seq_len = seq_len
+        self.lag = lag
+        self.device = device
+        self.cat_embeds = emb_dims + [seq_len, lag + 1, 2]  #
+        self.cat_n_embd = nn.ModuleList(
+            [nn.Embedding(emb_dim, d_model) for emb_dim in self.cat_embeds]
+        )
+
+    def forward(
+        self, x: Union[torch.Tensor, int], device: torch.device
+    ) -> torch.Tensor:
+        """All components of x are concatenated with 3 new variables for data augmentation, in the order:
+
+        - pos_seq: assign at each step its time-position
+        - pos_fut: assign at each step its future position. 0 if it is a past step
+        - is_fut: explicit for each step if it is a future(1) or past one(0)
+
+        Parameters
+        ----------
+            x: torch.Tensor
+                `[bs, seq_len, num_vars]`
+
+        Returns
+        ------
+            torch.Tensor:
+                `[bs, seq_len, num_vars+3, n_embd]`
+        """  # noqa: E501
+        if isinstance(x, int):
+            no_emb = True
+            B = x
+        else:
+            no_emb = False
+            B, _, _ = x.shape
+
+        pos_seq = self.get_pos_seq(bs=B).to(device)
+        pos_fut = self.get_pos_fut(bs=B).to(device)
+        is_fut = self.get_is_fut(bs=B).to(device)
+
+        if no_emb:
+            cat_vars = torch.cat((pos_seq, pos_fut, is_fut), dim=2)
+        else:
+            cat_vars = torch.cat((x, pos_seq, pos_fut, is_fut), dim=2)
+        cat_vars = cat_vars.long()
+        cat_n_embd = self.get_cat_n_embd(cat_vars)
+        return cat_n_embd
+
+    def get_pos_seq(self, bs):
+        pos_seq = torch.arange(0, self.seq_len)
+        pos_seq = pos_seq.repeat(bs, 1).unsqueeze(2).to(self.device)
+        return pos_seq
+
+    def get_pos_fut(self, bs):
+        pos_fut = torch.cat(
+            (
+                torch.zeros((self.seq_len - self.lag), dtype=torch.long),
+                torch.arange(1, self.lag + 1),
+            )
+        )
+        pos_fut = pos_fut.repeat(bs, 1).unsqueeze(2).to(self.device)
+        return pos_fut
+
+    def get_is_fut(self, bs):
+        is_fut = torch.cat(
+            (
+                torch.zeros((self.seq_len - self.lag), dtype=torch.long),
+                torch.ones((self.lag), dtype=torch.long),
+            )
+        )
+        is_fut = is_fut.repeat(bs, 1).unsqueeze(2).to(self.device)
+        return is_fut
+
+    def get_cat_n_embd(self, cat_vars):
+        cat_n_embd = torch.Tensor().to(cat_vars.device)
+        for index, layer in enumerate(self.cat_n_embd):
+            emb = layer(cat_vars[:, :, index])
+            cat_n_embd = torch.cat((cat_n_embd, emb.unsqueeze(2)), dim=2)
+        return cat_n_embd

pytorch_forecasting/models/tide/__init__.py

@@ -1,11 +1,8 @@
 """Tide model."""
 
 from pytorch_forecasting.models.tide._tide import TiDEModel
+from pytorch_forecasting.models.tide._tide_dsipts import TIDE, TIDE_pkg_v2
 from pytorch_forecasting.models.tide._tide_pkg import TiDEModel_pkg
 from pytorch_forecasting.models.tide.sub_modules import _TideModule
 
-__all__ = [
-    "_TideModule",
-    "TiDEModel",
-    "TiDEModel_pkg",
-]
+__all__ = ["_TideModule", "TiDEModel", "TiDEModel_pkg", "TIDE", "TIDE_pkg_v2"]

pytorch_forecasting/models/tide/_tide_dsipts/__init__.py

@@ -0,0 +1,6 @@
+"""DSIPTS Tide Implementation for V2"""
+
+from pytorch_forecasting.models.tide._tide_dsipts._tide_v2 import TIDE
+from pytorch_forecasting.models.tide._tide_dsipts._tide_v2_pkg import TIDE_pkg_v2
+
+__all__ = ["TIDE", "TIDE_pkg_v2"]