[ENH] [WIP] Standardize model output to 4d tensor

pytorch_forecasting/layers/decomposition/__init__.py

@@ -0,0 +1,9 @@
+"""
+Decomposition layers for PyTorch Forecasting.
+"""
+
+from pytorch_forecasting.layers.decomposition._series_decomp import SeriesDecomposition
+
+__all__ = [
+    "SeriesDecomposition",
+]

pytorch_forecasting/layers/decomposition/_series_decomp.py

@@ -0,0 +1,43 @@
+"""
+Series Decomposition Block for time series forecasting models.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pytorch_forecasting.layers.filter._moving_avg_filter import MovingAvg
+
+
+class SeriesDecomposition(nn.Module):
+    """
+    Series decomposition block from Autoformer.
+
+    Decomposes time series into trend and seasonal components using
+    moving average filtering.
+
+    Args:
+        kernel_size (int):
+            Size of the moving average kernel for trend extraction.
+    """
+
+    def __init__(self, kernel_size):
+        super().__init__()
+        self.moving_avg = MovingAvg(kernel_size, stride=1)
+
+    def forward(self, x):
+        """
+        Forward pass for series decomposition.
+
+        Args:
+            x (torch.Tensor):
+                Input time series tensor of shape (batch_size, seq_len, features).
+
+        Returns:
+            tuple:
+                - trend (torch.Tensor): Trend component of the time series.
+                - seasonal (torch.Tensor): Seasonal component of the time series.
+        """
+        trend = self.moving_avg(x)
+        seasonal = x - trend
+        return seasonal, trend

pytorch_forecasting/layers/filter/__init__.py

@@ -0,0 +1,9 @@
+"""
+Filtering layers for time series forecasting models.
+"""
+
+from pytorch_forecasting.layers.filter._moving_avg_filter import MovingAvg
+
+__all__ = [
+    "MovingAvg",
+]

pytorch_forecasting/layers/filter/_moving_avg_filter.py

@@ -0,0 +1,48 @@
+"""
+Moving Average Filter Block
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MovingAvg(nn.Module):
+    """
+    Moving Average block for smoothing and trend extraction from time series data.
+
+    A moving average is a smoothing technique that creates a series of average from
+    different subsets of a time series.
+
+    For example: Given a time series ``x = [x_1, x_2, ..., x_n]``, the moving average
+    with a kernel size of `k` calculates the average of each subset of `k` consecutive
+    elements, resulting in a new series of averages.
+
+    Args:
+        kernel_size (int):
+            Size of the moving average kernel.
+        stride (int):
+            Stride for the moving average operation, typically set to 1.
+    """
+
+    def __init__(self, kernel_size, stride):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.avg = nn.AvgPool1d(kernel_size, stride=stride, padding=0)
+
+    def forward(self, x):
+        if self.kernel_size % 2 == 0:
+            self.padding_left = self.kernel_size // 2 - 1
+            self.padding_right = self.kernel_size // 2
+        else:
+            self.padding_left = self.kernel_size // 2
+            self.padding_right = self.kernel_size // 2
+
+        front = x[:, 0:1, :].repeat(1, self.padding_left, 1)
+        end = x[:, -1:, :].repeat(1, self.padding_right, 1)
+
+        x_padded = torch.cat([front, x, end], dim=1)
+        x_transposed = x_padded.permute(0, 2, 1)
+        x_smoothed = self.avg(x_transposed)
+        x_out = x_smoothed.permute(0, 2, 1)
+        return x_out

pytorch_forecasting/models/base/_base_model_v2.py

@@ -294,3 +294,175 @@
                 prog_bar=True,
                 logger=True,
             )
+
+    def standardize_model_output(
+        self,
+        prediction: torch.Tensor,
+        expected_dims: tuple[int, Optional[int], Optional[int], Optional[int]] = (
+            None,
+            None,
+            None,
+            None,
+        ),  # noqa: E501
+    ) -> torch.Tensor:
+        """
+        Standardize model outputs to a 4-dimensional tensor, with shape
+        (batch_size, timesteps, num_features, last_dim).
+
+        Parameters
+        ----------
+        prediction : torch.Tensor
+            The raw prediction tensor from the model.
+            - Must be a torch.Tensor (in the future, also accept a list of tensors for
+              multi-target forecasting).
+            - Supported dims: 2D, 3D or 4D tensors.
+            - if 2D: (batch_size, timesteps) - univariate forecasting
+            - if 3D:
+                a) (batch_size, timesteps, n_targets) - multivariate forecasting
+                b) (batch_size, timesteps, last_dim) - univariate forecasting with quantiles or distribution.
+                c) (batch_size, timesteps, n_targets * last_dim) - multivariate
+              forecasting with quantiles, where features and quantiles are flattened in dim 2.
+            - if 4D: (batch_size, timesteps, n_targets, last_dim) - multivariate
+              forecasting with quantiles or distribution parameters.
+            - In the future, once multi-target forecasting is supported, this
+              will also accept a list of tensors, where each tensor inside the list
+              is treated as above.
+            - If anything apart from the above dimensions is provided, an error is raised.
+
+        expected_dims : tuple[int, Optional[int], Optional[int], Optional[int]], default=(None, None, None, None)
+        A tuple specifying the dimensions: (n_targets, batch_size, timesteps, last_dim)
+
+            n_targets : int
+                - Position 0: Number of target features
+                - Must be provided explicitly (cannot be None)
+                - Used for reshaping 2D and 3D tensors to 4D.
+
+            batch_size : Optional[int], default=None
+                - Position 1: Expected batch size
+                - When specified: Validates prediction.shape[0]
+                - When None: Uses actual tensor dimension
+
+            timesteps : Optional[int], default=None
+                - Position 2: Expected number of timesteps
+                - When specified: Validates prediction.shape[1]
+                - When None: Uses actual tensor dimension
+
+            last_dim : Optional[int], default=None
+                - Position 3: Size of the last dimension.
+                - Common use case - quantile, sample, distribution params.
+                - When it is specified, it is used to directly reshape.
+                - When None and model uses QuantileLoss: It is set to the number of quantiles
+                - When None and no quantile information is available: It defaults to 1.
+                - If required, this can be extended to handle other cases where the last_dim is None
+                but its value can be inferred from the loss function or model configuration (apart from
+                the existing QuantileLoss case, of course).
+        Returns
+        -------
+        torch.Tensor
+            The standardized prediction tensor with shape (batch_size, timesteps, n_targets, last_dim).
+            The prediction tensor is obtained by reshaping the input tensor. There are
+            several cases to consider:
+
+            - If the input tensor is 2D, it is reshaped to (batch_size, timesteps, 1, 1).
+            - If the input tensor is 3D, it is reshaped to (batch_size, timesteps, 1, 1) for a
+            multivariate single-target forecast, or (batch_size, timesteps, 1, last_dim) for a univariate quantile forecast.
+            - If the input tensor is 4D, it is assumed to be in the shape
+            (batch_size, timesteps, n_targets, last_dim) or (batch_size, timesteps, last_dim, n_targets).
+        Notes
+        -----
+        [1] The fourth dimension (last_dim) commonly represents:
+
+            * Quantiles: For quantile regression (e.g., 0.1, 0.5, 0.9)
+            * Distribution parameters: For parametric forecasts (e.g., mean, variance)
+            * Samples: For sample-based uncertainty estimates
+
+            The current implementation assumes the most common case of quantile forecasts
+            when automatically inferring this dimension from the loss function,
+            but any value can be explicitly provided. A fallback of 1 is used in case where
+            no information is available on ``last_dim``.
+
+        [2] This can currently handle situations where a single target is used
+            either in a univariate or multivariate situation. In case of multi-target
+            forecasting, where each target has its own loss function, a list of tensors is
+            returned, where each tensor corresponds to a target. This requires some change
+            to the existing code.
+        """  # noqa: E501
+
+        n_targets, batch_size, timesteps, last_dim = expected_dims
+
+        if not isinstance(prediction, torch.Tensor):
+            raise TypeError(
+                f"Expected prediction to be a torch.Tensor, but got {type(prediction)}"
+            )
+
+        if n_targets is None:
+            raise ValueError(
+                "Expected n_targets to be a positive integer, but got `None`."
+            )
+
+        if last_dim is None:
+            if hasattr(self.loss, "quantiles") and self.loss.quantiles is not None:
+                last_dim = len(self.loss.quantiles)  # Quantile regression case
+            # we can add more cases here in the future, where we refer to the specific
+            # loss function to determine the last dimension. For now we are sticking
+            # to the quantile regression case.
+            else:
+                last_dim = 1
+
+        if batch_size is not None:
+            if prediction.shape[0] != batch_size:
+                raise ValueError(
+                    f"Expected batch size {batch_size}, but got {prediction.shape[0]}."
+                )
+
+        if timesteps is not None:
+            if prediction.shape[1] != timesteps:
+                raise ValueError(
+                    f"Expected timesteps {timesteps}, but got {prediction.shape[1]}."
+                )
+
+        if prediction.ndim == 2:
+            # reshape to (batch_size, timsteps, 1, 1)
+            prediction = prediction.unsqueeze(-1).unsqueeze(-1)
+
+        elif prediction.ndim == 3:
+            if prediction.shape[2] == n_targets:
+                # reshape to (batch_size, timesteps, n_targets, 1)
+                prediction = prediction.unsqueeze(-1)
+            elif prediction.shape[2] == last_dim:
+                # reshape to (batch_size, timesteps, 1, last_dim)
+                prediction = prediction.unsqueeze(2)
+            elif prediction.shape[2] == n_targets * last_dim:
+                # multivariate forecast with quantiles
+                # where features and quantiles are flattened in dim 2.
+                # reshape to (batch_size, timesteps, n_targets, last_dim)
+                prediction = prediction.reshape(
+                    prediction.shape[0], prediction.shape[1], n_targets, last_dim
+                )
+            else:
+                # reshape to (batch_size, timesteps, n_targets, last_dim)
+                prediction = prediction.unsqueeze(-1)
+
+        elif prediction.ndim == 4:
+            # assuming only a single case where n_targets and last_dim are swapped.
+            if prediction.shape[2] == last_dim and prediction.shape[3] == n_targets:
+                # reshape to (batch_size, timesteps, n_targets, last_dim)
+                warn(
+                    "Prediction tensor has shape (batch_size, timesteps, last_dim, n_targets). "  # noqa: E501
+                    "This is not the expected shape. Transposing the last two dimensions."  # noqa: E501
+                )
+                prediction = prediction.permute(0, 1, 3, 2)
+
+        else:
+            raise ValueError(
+                f"Expected prediction tensor to have 2, 3, or 4 dimensions, "
+                f"but got {prediction.ndim} dimensions."
+            )
+
+        # final check to ensure the output is 4D
+        if prediction.ndim != 4:
+            raise ValueError(
+                f"Failed to standardize output to 4D tensor. Current shape: {prediction.shape}"  # noqa: E501
+            )
+
+        return prediction

pytorch_forecasting/models/dlinear/__init__.py

@@ -0,0 +1,10 @@
+"""
+Decomposition-Linear model for time series forecasting.
+"""
+
+from pytorch_forecasting.models.dlinear._dlinear_pkg_v2 import DLinearModel_pkg_v2
+from pytorch_forecasting.models.dlinear._dlinear_v2 import DLinearModel
+
+__all__ = [
+    "DLinearModel" "DLinearModel_pkg_v2",
+]