[ENH] Implement NBeatsKAN in v2 interface

Author: Dev10-sys

pytorch_forecasting/models/__init__.py

@@ -11,7 +11,12 @@
 from pytorch_forecasting.models.baseline import Baseline
 from pytorch_forecasting.models.deepar import DeepAR
 from pytorch_forecasting.models.mlp import DecoderMLP
-from pytorch_forecasting.models.nbeats import NBeats, NBeatsKAN
+from pytorch_forecasting.models.nbeats import (
+    NBeats,
+    NBeatsKAN,
+    NBeatsKAN_pkg_v2,
+    NBeatsKAN_v2,
+)
 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
@@ -26,6 +31,8 @@
 __all__ = [
     "NBeats",
     "NBeatsKAN",
+    "NBeatsKAN_v2",
+    "NBeatsKAN_pkg_v2",
     "NHiTS",
     "TemporalFusionTransformer",
     "RecurrentNetwork",

pytorch_forecasting/models/nbeats/__init__.py

@@ -16,10 +16,14 @@
 from pytorch_forecasting.models.nbeats._nbeats_pkg import NBeats_pkg
 from pytorch_forecasting.models.nbeats._nbeatskan import NBeatsKAN
 from pytorch_forecasting.models.nbeats._nbeatskan_pkg import NBeatsKAN_pkg
+from pytorch_forecasting.models.nbeats._nbeatskan_pkg_v2 import NBeatsKAN_pkg_v2
+from pytorch_forecasting.models.nbeats._nbeatskan_v2 import NBeatsKAN_v2
 
 __all__ = [
     "NBeats",
     "NBeatsKAN",
+    "NBeatsKAN_v2",
+    "NBeatsKAN_pkg_v2",
     "NBeats_pkg",
     "NBeatsKAN_pkg",
     "NBEATSGenericBlock",

pytorch_forecasting/models/nbeats/_nbeatskan_pkg_v2.py

@@ -0,0 +1,62 @@
+"""NBeatsKAN v2 package container."""
+
+from typing import Any
+
+from pytorch_forecasting.base._base_pkg import Base_pkg
+
+
+class NBeatsKAN_pkg_v2(Base_pkg):
+    """NBeatsKAN v2 package container."""
+
+    _tags: dict[str, Any] = {
+        "info:name": "NBeatsKAN",
+        "info:compute": 2,
+        "info:y_type": ["numeric"],
+        "authors": ["Dev10-sys", "kalpanagola9897"],
+        "capability:exogenous": False,
+        "capability:multivariate": False,
+        "capability:pred_int": True,
+        "capability:flexible_history_length": False,
+        "capability:cold_start": False,
+        "python_dependencies": [],
+    }
+
+    @classmethod
+    def get_cls(cls):
+        """Get the model class."""
+        from pytorch_forecasting.models.nbeats._nbeatskan_v2 import NBeatsKAN_v2
+
+        return NBeatsKAN_v2
+
+    @classmethod
+    def get_datamodule_cls(cls):
+        """Get the compatible DataModule class."""
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
+
+        return EncoderDecoderTimeSeriesDataModule
+
+    @classmethod
+    def get_test_train_params(cls) -> list[dict[str, Any]]:
+        """Return testing parameter settings for trainer fixtures."""
+        from pytorch_forecasting.metrics import MAE
+
+        params: list[dict[str, Any]] = [
+            {},
+            {"stack_types": ["generic"], "num_blocks": [1], "widths": [32]},
+            {"loss": MAE()},
+        ]
+
+        default_dm_cfg: dict[str, Any] = {
+            "max_encoder_length": 8,
+            "max_prediction_length": 3,
+        }
+
+        for param in params:
+            current_dm_cfg = param.get("datamodule_cfg", {})
+            if isinstance(current_dm_cfg, dict):
+                default_dm_cfg.update(current_dm_cfg)
+            param["datamodule_cfg"] = default_dm_cfg.copy()
+
+        return params

pytorch_forecasting/models/nbeats/_nbeatskan_v2.py

@@ -0,0 +1,244 @@
+"""NBeatsKAN model for PyTorch Forecasting v2."""
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch import nn
+from torch.optim import Optimizer
+
+from pytorch_forecasting.layers._nbeats._blocks import (
+    NBEATSGenericBlockKAN,
+    NBEATSSeasonalBlockKAN,
+    NBEATSTrendBlockKAN,
+)
+from pytorch_forecasting.metrics import MAE, Metric
+from pytorch_forecasting.models.base._base_model_v2 import BaseModel
+
+
+class NBeatsKAN_v2(BaseModel):
+    """N-BEATS model with Kolmogorov-Arnold Network (KAN) spline layers for v2.
+
+    Parameters
+    ----------
+    loss : Metric, default=MAE()
+        Loss metric used for training and evaluation.
+    stack_types : list of str, optional
+        List of stack types: ``"generic"``, ``"trend"``, or ``"seasonality"``.
+    num_blocks : list of int, optional
+        Number of blocks per stack.
+    num_block_layers : list of int, optional
+        Number of KAN layers per block.
+    widths : list of int, optional
+        Widths of layers in blocks.
+    sharing : list of bool, optional
+        Whether blocks share weights per stack.
+    expansion_coefficient_lengths : list of int, optional
+        Expansion lengths or polynomial degrees per stack.
+    dropout : float, default=0.1
+        Dropout rate.
+    logging_metrics : list of nn.Module, optional
+        Logged evaluation metrics.
+    optimizer : Optimizer or str, default="adam"
+        Optimizer used for training.
+    optimizer_params : dict, optional
+        Optimizer parameters.
+    lr_scheduler : str, optional
+        Learning rate scheduler name.
+    lr_scheduler_params : dict, optional
+        Parameters for the learning rate scheduler.
+    num : int, default=5
+        KAN grid intervals.
+    k : int, default=3
+        KAN spline polynomial order.
+    noise_scale : float, default=0.5
+        KAN noise scale at initialization.
+    scale_base_mu : float, default=0.0
+        KAN base scale mean.
+    scale_base_sigma : float, default=1.0
+        KAN base scale std.
+    scale_sp : float, default=1.0
+        KAN spline scale.
+    base_fun : Callable, optional
+        KAN residual base activation function.
+    grid_eps : float, default=0.02
+        KAN grid interpolation parameter.
+    grid_range : list of int, optional
+        KAN grid range boundaries.
+    sp_trainable : bool, default=True
+        Whether spline scale is trainable.
+    sb_trainable : bool, default=True
+        Whether base scale is trainable.
+    sparse_init : bool, default=False
+        Whether sparse initialization is used.
+    metadata : dict, optional
+        Metadata from DataModule.
+    """
+
+    @classmethod
+    def _pkg(cls):
+        """Package container for the model."""
+        from pytorch_forecasting.models.nbeats._nbeatskan_pkg_v2 import (
+            NBeatsKAN_pkg_v2,
+        )
+
+        return NBeatsKAN_pkg_v2
+
+    def __init__(
+        self,
+        loss: Metric = MAE(),
+        stack_types: list[str] | None = None,
+        num_blocks: list[int] | None = None,
+        num_block_layers: list[int] | None = None,
+        widths: list[int] | None = None,
+        sharing: list[bool] | None = None,
+        expansion_coefficient_lengths: list[int] | None = None,
+        dropout: float = 0.1,
+        logging_metrics: list[nn.Module] | None = None,
+        optimizer: Optimizer | str | None = "adam",
+        optimizer_params: dict | None = None,
+        lr_scheduler: str | None = None,
+        lr_scheduler_params: dict | None = None,
+        num: int = 5,
+        k: int = 3,
+        noise_scale: float = 0.5,
+        scale_base_mu: float = 0.0,
+        scale_base_sigma: float = 1.0,
+        scale_sp: float = 1.0,
+        base_fun: Callable | None = None,
+        grid_eps: float = 0.02,
+        grid_range: list[int] | None = None,
+        sp_trainable: bool = True,
+        sb_trainable: bool = True,
+        sparse_init: bool = False,
+        metadata: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            loss=loss,
+            logging_metrics=logging_metrics,
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            lr_scheduler=lr_scheduler,
+            lr_scheduler_params=lr_scheduler_params,
+        )
+        self.save_hyperparameters(
+            ignore=["loss", "logging_metrics", "metadata", "base_fun"]
+        )
+        self.metadata = metadata or {}
+
+        self.max_encoder_length = self.metadata.get("max_encoder_length", 10)
+        self.max_prediction_length = self.metadata.get("max_prediction_length", 1)
+
+        if base_fun is None:
+            base_fun = nn.SiLU()
+        if grid_range is None:
+            grid_range = [-1, 1]
+        if expansion_coefficient_lengths is None:
+            expansion_coefficient_lengths = [3, 7]
+        if sharing is None:
+            sharing = [True, True]
+        if widths is None:
+            widths = [32, 512]
+        if num_block_layers is None:
+            num_block_layers = [3, 3]
+        if num_blocks is None:
+            num_blocks = [3, 3]
+        if stack_types is None:
+            stack_types = ["trend", "seasonality"]
+
+        kan_params = {
+            "num": num,
+            "k": k,
+            "noise_scale": noise_scale,
+            "scale_base_mu": scale_base_mu,
+            "scale_base_sigma": scale_base_sigma,
+            "scale_sp": scale_sp,
+            "base_fun": base_fun,
+            "grid_eps": grid_eps,
+            "grid_range": grid_range,
+            "sp_trainable": sp_trainable,
+            "sb_trainable": sb_trainable,
+            "sparse_init": sparse_init,
+        }
+
+        self.net_blocks = nn.ModuleList()
+        for stack_id, stack_type in enumerate(stack_types):
+            for _ in range(num_blocks[stack_id]):
+                net_block: nn.Module
+                if stack_type == "generic":
+                    net_block = NBEATSGenericBlockKAN(
+                        units=widths[stack_id],
+                        thetas_dim=expansion_coefficient_lengths[stack_id],
+                        num_block_layers=num_block_layers[stack_id],
+                        backcast_length=self.max_encoder_length,
+                        forecast_length=self.max_prediction_length,
+                        dropout=dropout,
+                        **kan_params,
+                    )
+                elif stack_type == "seasonality":
+                    net_block = NBEATSSeasonalBlockKAN(
+                        units=widths[stack_id],
+                        thetas_dim=expansion_coefficient_lengths[stack_id],
+                        num_block_layers=num_block_layers[stack_id],
+                        backcast_length=self.max_encoder_length,
+                        forecast_length=self.max_prediction_length,
+                        nb_harmonics=None,  # type: ignore[arg-type]
+                        min_period=expansion_coefficient_lengths[stack_id],
+                        dropout=dropout,
+                        **kan_params,
+                    )
+
+                elif stack_type == "trend":
+                    net_block = NBEATSTrendBlockKAN(
+                        units=widths[stack_id],
+                        thetas_dim=expansion_coefficient_lengths[stack_id],
+                        num_block_layers=num_block_layers[stack_id],
+                        backcast_length=self.max_encoder_length,
+                        forecast_length=self.max_prediction_length,
+                        dropout=dropout,
+                        **kan_params,
+                    )
+                else:
+                    raise ValueError(f"Unknown stack_type: {stack_type}")
+
+                self.net_blocks.append(net_block)
+
+    def forward(
+        self,
+        x: dict[str, torch.Tensor],
+    ) -> dict[str, torch.Tensor]:
+        """Forward pass for NBeatsKAN_v2.
+
+        Parameters
+        ----------
+        x : dict[str, torch.Tensor]
+            Input dictionary containing ``target_past``.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            Dictionary containing predicted output tensor under key ``prediction``.
+        """
+        target_past = x.get("target_past")
+        if target_past is None:
+            raise KeyError("Input dictionary must contain key 'target_past'.")
+
+        if target_past.ndim == 3 and target_past.size(-1) == 1:
+            target_past = target_past.squeeze(-1)
+
+        backcast = target_past
+        forecast = torch.zeros(
+            backcast.size(0),
+            self.max_prediction_length,
+            device=backcast.device,
+            dtype=backcast.dtype,
+        )
+
+        for block in self.net_blocks:
+            b, f = block(backcast)
+            backcast = backcast - b
+            forecast = forecast + f
+
+        prediction = forecast.unsqueeze(-1)
+        return {"prediction": prediction}