[ENH] softs_v2 Model added

pytorch_forecasting/layers/_blocks/__init__.py

@@ -1,3 +1,11 @@
 from pytorch_forecasting.layers._blocks._residual_block_dsipts import ResidualBlock
+from pytorch_forecasting.layers._blocks._softs_block import (
+    SoftsEncoderLayer,
+    STADModule,
+)
 
-__all__ = ["ResidualBlock"]
+__all__ = [
+    "ResidualBlock",
+    "STADModule",
+    "SoftsEncoderLayer",
+]

pytorch_forecasting/layers/_blocks/_softs_block.py

@@ -0,0 +1,114 @@
+"""
+SOFTS Blocks for Star Aggregate-Dispatch Network.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class STADModule(nn.Module):
+    """
+    Star Aggregate-Dispatch (STAD) Module for capturing inter-series dependencies.
+
+    Uses a star-topology to aggregate all channels into a central node,
+    process it via an MLP, and dispatch back — achieving O(C) cross-channel
+    mixing instead of O(C²) self-attention.
+
+    Parameters
+    ----------
+    d_model : int
+        Embedding dimension per channel per time step.
+    d_core : int
+        Dimension of the central star node (information bottleneck).
+    dropout : float, default=0.0
+        Dropout probability inside the channel-mixing MLP.
+    """
+
+    def __init__(self, d_model: int, d_core: int, dropout: float = 0.0):
+        super().__init__()
+        self.channel_mixing = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_model, d_model),
+        )
+        self.gen_weight = nn.Linear(d_model, d_core)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Aggregate channel features into a star node and dispatch back.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Shape ``(batch_size, n_channels, seq_len, d_model)``.
+
+        Returns
+        -------
+        torch.Tensor
+            Same shape as input, enriched with cross-channel context.
+        """
+
+        B, C, L, D = x.shape
+
+        w = self.gen_weight(x).mean(dim=2)
+        w = torch.softmax(w, dim=1)
+
+        x_pooled = x.mean(dim=2)
+        core_node = torch.einsum("bcd,bce->bed", x_pooled, w)
+        core_node = self.channel_mixing(core_node)
+        dispatch_out = torch.einsum("bed,bce->bcd", core_node, w)
+
+        dispatch_out = dispatch_out.unsqueeze(2).repeat(1, 1, L, 1)
+        return x + dispatch_out
+
+
+class SoftsEncoderLayer(nn.Module):
+    """
+    Single encoder layer for SOFTS, combining STAD and a Feed-Forward Network.
+
+    Applies Pre-LayerNorm STAD (cross-channel) then FFN (within-channel)
+    with residual connections, following the Pre-LN Transformer convention.
+
+    Parameters
+    ----------
+    d_model : int
+        Embedding dimension per channel per time step.
+    d_core : int
+        Dimension of the central star node in the STAD sub-layer.
+    d_ff : int
+        Hidden dimension of the feed-forward network (typically 4 x d_model).
+    dropout : float, default=0.0
+        Dropout probability applied after the STAD and FFN sub-layers.
+    """
+
+    def __init__(self, d_model: int, d_core: int, d_ff: int, dropout: float = 0.0):
+        super().__init__()
+        self.stad = STADModule(d_model=d_model, d_core=d_core, dropout=dropout)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply one SOFTS encoder layer: STAD sub-layer then FFN sub-layer.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape ``(batch_size, n_channels, seq_len, d_model)``.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape ``(batch_size, n_channels, seq_len, d_model)``.
+        """
+        x = x + self.dropout(self.stad(self.norm1(x)))
+        x = x + self.dropout(self.ffn(self.norm2(x)))
+
+        return x

pytorch_forecasting/models/__init__.py

@@ -15,6 +15,7 @@
 from pytorch_forecasting.models.nhits import NHiTS
 from pytorch_forecasting.models.nn import GRU, LSTM, MultiEmbedding, get_rnn
 from pytorch_forecasting.models.rnn import RecurrentNetwork
+from pytorch_forecasting.models.softs import Softs, Softs_pkg_v2
 from pytorch_forecasting.models.temporal_fusion_transformer import (
     TemporalFusionTransformer,
 )
@@ -42,4 +43,6 @@
     "TiDEModel",
     "TimeXer",
     "xLSTMTime",
+    "Softs",
+    "Softs_pkg_v2",
 ]

pytorch_forecasting/models/softs/__init__.py

@@ -0,0 +1,8 @@
+"""
+SOFTS Model for Multivariate Time Series Forecasting.
+"""
+
+from pytorch_forecasting.models.softs._softs_pkg_v2 import Softs_pkg_v2
+from pytorch_forecasting.models.softs._softs_v2 import Softs
+
+__all__ = ["Softs", "Softs_pkg_v2"]

pytorch_forecasting/models/softs/_softs_pkg_v2.py

@@ -0,0 +1,87 @@
+"""
+Packages container for SOFTS model.
+"""
+
+from pytorch_forecasting.base._base_pkg import Base_pkg
+
+
+class Softs_pkg_v2(Base_pkg):
+    """
+    SOFTS package container.
+    Reference : https://arxiv.org/abs/2404.14197
+    """
+
+    _tags = {
+        "info:name": "SOFTS",
+        "info:y_type": ["numeric"],
+        "info:compute": 2,
+        "authors": ["Secilia-Cxy", "Muhammad-Rebaal"],
+        "capability:exogenous": True,
+        "capability:multivariate": True,
+        "capability:pred_int": True,
+        "capability:flexible_history_length": True,
+        "capability:cold_start": False,
+    }
+
+    @classmethod
+    def get_cls(cls):
+        """Get model class."""
+        from pytorch_forecasting.models.softs._softs_v2 import Softs
+
+        return Softs
+
+    @classmethod
+    def get_datamodule_cls(cls):
+        """Get the underlying DataModule class."""
+        from pytorch_forecasting.data.data_module import TslibDataModule
+
+        return TslibDataModule
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Return testing parameter settings for the trainer.
+
+        Returns
+        -------
+        list of dict
+            Each dict is a valid set of constructor arguments for ``Softs``.
+            The key ``datamodule_cfg`` is passed to the DataModule, not the model.
+        """
+        from pytorch_forecasting.metrics import SMAPE, QuantileLoss
+
+        params = [
+            {},
+            dict(hidden_size=64, d_core=64, d_ff=256, n_layers=1),
+            dict(hidden_size=128, n_layers=1, use_revin=False),
+            dict(
+                hidden_size=64,
+                n_layers=1,
+                loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
+            ),
+            dict(
+                hidden_size=64,
+                n_layers=1,
+                use_revin=False,
+                loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
+            ),
+            dict(hidden_size=64, dropout=0.0, n_layers=1),
+            dict(datamodule_cfg=dict(context_length=16, prediction_length=4)),
+            dict(
+                optimizer="adamw",
+                lr_scheduler="cosine_annealing",
+                lr_scheduler_params={"T_max": 5},
+            ),
+            dict(
+                optimizer="adagrad",
+                optimizer_params={"lr": 1e-3},
+            ),
+            dict(hidden_size=64, n_layers=1, logging_metrics=[SMAPE()]),
+        ]
+
+        default_dm_cfg = {"context_length": 8, "prediction_length": 2}
+
+        for param in params:
+            current_dm_cfg = param.get("datamodule_cfg", {})
+            param["datamodule_cfg"] = {**default_dm_cfg, **current_dm_cfg}
+
+        return params