diff --git a/docs/source/m_layer_v2.rst b/docs/source/m_layer_v2.rst index 80b447f95..dd16256dc 100644 --- a/docs/source/m_layer_v2.rst +++ b/docs/source/m_layer_v2.rst @@ -52,3 +52,6 @@ See the detailed API documentation for the V2 base classes and specific model im models.softs._softs_v2.SOFTS models.scinet._scinet_v2.SCINet_v2 models.patch_tst._patch_tst_v2.PatchTST_v2 + models.nbeats._nbeats_adapter_v2.NBeatsAdapterV2 + models.nbeats._nbeats_v2.NBeats + models.nbeats._nbeatskan_v2.NBeatsKAN_v2 diff --git a/docs/source/pkg_v2.rst b/docs/source/pkg_v2.rst index ef296de6b..a077def56 100644 --- a/docs/source/pkg_v2.rst +++ b/docs/source/pkg_v2.rst @@ -104,3 +104,5 @@ See the detailed API documentation for the available V2 Package classes below: models.softs._softs_pkg_v2.SOFTS_pkg_v2 models.scinet._scinet_pkg_v2.SCINet_pkg_v2 models.patch_tst._patch_tst_pkg_v2.PatchTST_pkg_v2 + models.nbeats._nbeats_pkg_v2.NBeats_pkg_v2 + models.nbeats._nbeatskan_pkg_v2.NBeatsKAN_pkg_v2 diff --git a/pytorch_forecasting/models/__init__.py b/pytorch_forecasting/models/__init__.py index 2c9f71c1b..a760d18e6 100644 --- a/pytorch_forecasting/models/__init__.py +++ b/pytorch_forecasting/models/__init__.py @@ -11,7 +11,15 @@ 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, + NBeats_pkg_v2, + NBeats_v2, + NBeatsAdapterV2, + 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.patch_tst import ( @@ -32,7 +40,12 @@ __all__ = [ "NBeats", + "NBeats_v2", + "NBeats_pkg_v2", + "NBeatsAdapterV2", "NBeatsKAN", + "NBeatsKAN_v2", + "NBeatsKAN_pkg_v2", "NHiTS", "PatchTST", "PatchTST_v2", diff --git a/pytorch_forecasting/models/nbeats/__init__.py b/pytorch_forecasting/models/nbeats/__init__.py index 5377ff67f..58d246f4b 100644 --- a/pytorch_forecasting/models/nbeats/__init__.py +++ b/pytorch_forecasting/models/nbeats/__init__.py @@ -13,18 +13,28 @@ from pytorch_forecasting.models.nbeats._grid_callback import GridUpdateCallback from pytorch_forecasting.models.nbeats._nbeats import NBeats from pytorch_forecasting.models.nbeats._nbeats_adapter import NBeatsAdapter +from pytorch_forecasting.models.nbeats._nbeats_adapter_v2 import NBeatsAdapterV2 from pytorch_forecasting.models.nbeats._nbeats_pkg import NBeats_pkg +from pytorch_forecasting.models.nbeats._nbeats_pkg_v2 import NBeats_pkg_v2 +from pytorch_forecasting.models.nbeats._nbeats_v2 import NBeats_v2 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", + "NBeats_v2", + "NBeats_pkg_v2", "NBeatsKAN", + "NBeatsKAN_v2", + "NBeatsKAN_pkg_v2", "NBeats_pkg", "NBeatsKAN_pkg", "NBEATSGenericBlock", "NBEATSSeasonalBlock", "NBEATSTrendBlock", "NBeatsAdapter", + "NBeatsAdapterV2", "GridUpdateCallback", ] diff --git a/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py new file mode 100644 index 000000000..1cbde933e --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeats_adapter_v2.py @@ -0,0 +1,162 @@ +"""Shared N-Beats adapter for pytorch-forecasting v2.""" + +from typing import Any + +import torch +from torch import nn +from torch.optim import Optimizer + +from pytorch_forecasting.layers._nbeats._blocks import ( + NBEATSSeasonalBlock, + NBEATSTrendBlock, + SeasonalMixin, + TrendMixin, +) +from pytorch_forecasting.metrics import Metric +from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel + + +class NBeatsAdapterV2(TslibBaseModel): + """Shared forward / training helpers for NBeats and NBeatsKAN (v2).""" + + def __init__( + self, + loss: Metric, + 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, + metadata: dict | None = None, + backcast_loss_ratio: float = 0.0, + **kwargs: Any, + ): + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + metadata=metadata, + ) + self.backcast_loss_ratio = backcast_loss_ratio + + def _target_from_batch(self, x: dict[str, torch.Tensor]) -> torch.Tensor: + """Extract univariate target history. + + v1 used ``x["encoder_cont"][..., 0]``. v2 tslib batches keep the target + in ``history_target``. + """ + target = x["history_target"] + if target.ndim == 3: + target = target[..., 0] + return target + + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Pass forward of network. + + Network steps match v1 ``NBeatsAdapter.forward``; only input assembly + and output packaging differ for the v2 API. + """ + # --- v2 batch adapter (v1: target = x["encoder_cont"][..., 0]) --- + target = self._target_from_batch(x) + + # --- same as v1 from here --- + timesteps = self.context_length + self.prediction_length + generic_forecast = [ + torch.zeros( + (target.size(0), timesteps), dtype=torch.float32, device=self.device + ) + ] + trend_forecast = [ + torch.zeros( + (target.size(0), timesteps), dtype=torch.float32, device=self.device + ) + ] + seasonal_forecast = [ + torch.zeros( + (target.size(0), timesteps), dtype=torch.float32, device=self.device + ) + ] + forecast = torch.zeros( + (target.size(0), self.prediction_length), + dtype=torch.float32, + device=self.device, + ) + + backcast = target # initialize backcast + for i, block in enumerate(self.net_blocks): + # evaluate block + backcast_block, forecast_block = block(backcast) + + # add for interpretation + full = torch.cat([backcast_block.detach(), forecast_block.detach()], dim=1) + if isinstance(block, (NBEATSTrendBlock, TrendMixin)): + trend_forecast.append(full) + elif isinstance(block, (NBEATSSeasonalBlock, SeasonalMixin)): + seasonal_forecast.append(full) + else: + generic_forecast.append(full) + + # update backcast and forecast + backcast = ( + backcast - backcast_block + ) # do not use backcast -= backcast_block as this signifies an inline operation # noqa: E501 + forecast = forecast + forecast_block + + prediction = forecast.unsqueeze(-1) + backcast_out = (target - backcast).unsqueeze(-1) + trend = torch.stack(trend_forecast, dim=0).sum(0).unsqueeze(-1) + seasonality = torch.stack(seasonal_forecast, dim=0).sum(0).unsqueeze(-1) + generic = torch.stack(generic_forecast, dim=0).sum(0).unsqueeze(-1) + + # v1 applied transform_output via BaseModel; v2 tslib does so when scales exist + if "target_scale" in x: + prediction = self.transform_output(prediction, x["target_scale"]) + backcast_out = self.transform_output(backcast_out, x["target_scale"]) + trend = self.transform_output(trend, x["target_scale"]) + seasonality = self.transform_output(seasonality, x["target_scale"]) + generic = self.transform_output(generic, x["target_scale"]) + + # v1: to_network_output(...); v2: plain dict + return { + "prediction": prediction, + "backcast": backcast_out, + "trend": trend, + "seasonality": seasonality, + "generic": generic, + } + + def training_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> dict[str, torch.Tensor]: + """Training step with optional backcast loss (v1 ``step`` parity).""" + x, y = batch + out = self(x) + y_hat = out["prediction"] + loss = self.loss(y_hat, y) + + if self.backcast_loss_ratio > 0: + backcast = out["backcast"].squeeze(-1) + encoder_target = self._target_from_batch(x) + + backcast_weight = ( + self.backcast_loss_ratio + * self.prediction_length + / max(self.context_length, 1) + ) + backcast_weight = backcast_weight / (backcast_weight + 1) + forecast_weight = 1 - backcast_weight + + # Compute backcast term directly (avoid Metric.update state / shape quirks). + # v1 used self.loss(backcast, encoder_target); v2 BaseModel losses are + # wired for forecast horizon shapes only. + backcast_loss = (backcast - encoder_target).abs().mean() * backcast_weight + loss = loss * forecast_weight + backcast_loss + + self.log( + "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="train") + return {"loss": loss} diff --git a/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py new file mode 100644 index 000000000..c83cfe3f4 --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeats_pkg_v2.py @@ -0,0 +1,98 @@ +"""NBeats v2 package container.""" + +from pytorch_forecasting.base._base_pkg import Base_pkg + + +class NBeats_pkg_v2(Base_pkg): + """NBeats v2 package container.""" + + _tags = { + "info:name": "NBeats", + "info:compute": 1, + "info:y_type": ["numeric"], + "authors": [ + "dmitri-carpov", # paper author + "jdb78", # for v1 + "Faakhir30", + ], + "capability:exogenous": False, + "capability:multivariate": False, + "capability:pred_int": False, + "capability:flexible_history_length": False, + "capability:cold_start": False, + } + + @classmethod + def get_cls(cls): + """Get model class.""" + from pytorch_forecasting.models.nbeats._nbeats_v2 import NBeats + + return NBeats + + @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.""" + from pytorch_forecasting.metrics import MAE, MAPE, SMAPE + + params = [ + { + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "backcast_loss_ratio": 1.0, + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "stack_types": ["generic"], + "num_blocks": [1], + "num_block_layers": [2], + "widths": [16], + "expansion_coefficient_lengths": [8], + "sharing": [False], + }, + { + "loss": MAE(), + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "loss": MAPE(), + "logging_metrics": [SMAPE()], + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "optimizer": "adamw", + "lr_scheduler": "cosine_annealing", + "lr_scheduler_params": {"T_max": 5}, + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + ] + + default_dm_cfg = { + "context_length": 8, + "prediction_length": 3, + "add_relative_time_idx": False, + } + + for param in params: + current_dm_cfg = param.get("datamodule_cfg", {}) + default_dm_cfg.update(current_dm_cfg) + param["datamodule_cfg"] = default_dm_cfg.copy() + + return params diff --git a/pytorch_forecasting/models/nbeats/_nbeats_v2.py b/pytorch_forecasting/models/nbeats/_nbeats_v2.py new file mode 100644 index 000000000..f01d1b918 --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeats_v2.py @@ -0,0 +1,133 @@ +""" +N-Beats model for pytorch-forecasting v2 (no covariates). +""" + +from typing import Any + +from torch import nn +from torch.optim import Optimizer + +from pytorch_forecasting.layers._nbeats._blocks import ( + NBEATSGenericBlock, + NBEATSSeasonalBlock, + NBEATSTrendBlock, +) +from pytorch_forecasting.metrics import MAE, MAPE, RMSE, SMAPE, Metric +from pytorch_forecasting.models.nbeats._nbeats_adapter_v2 import NBeatsAdapterV2 + + +class NBeats(NBeatsAdapterV2): + """ + N-BEATS for pytorch-forecasting v2. + + Based on + `N-BEATS: Neural basis expansion analysis for interpretable time series + forecasting `_. + + Network construction matches the v1 ``NBeats`` class; ``context_length`` / + ``prediction_length`` come from datamodule ``metadata`` instead of + ``from_dataset``. + """ + + @classmethod + def _pkg(cls): + """Package for the model.""" + from pytorch_forecasting.models.nbeats._nbeats_pkg_v2 import NBeats_pkg_v2 + + return NBeats_pkg_v2 + + def __init__( + self, + loss: Metric, + 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, + backcast_loss_ratio: float = 0.0, + 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, + metadata: dict | None = None, + **kwargs: Any, + ): + 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"] + if logging_metrics is None: + logging_metrics = [SMAPE(), MAE(), RMSE(), MAPE()] + + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + metadata=metadata, + backcast_loss_ratio=backcast_loss_ratio, + ) + self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"]) + + self.stack_types = stack_types + self.num_blocks = num_blocks + self.num_block_layers = num_block_layers + self.widths = widths + self.sharing = sharing + self.expansion_coefficient_lengths = expansion_coefficient_lengths + self.dropout = dropout + + self._init_network() + + def _init_network(self): + """Build N-BEATS stacks (same block wiring as v1).""" + self.net_blocks = nn.ModuleList() + for stack_id, stack_type in enumerate(self.stack_types): + for _ in range(self.num_blocks[stack_id]): + if stack_type == "generic": + net_block = NBEATSGenericBlock( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + dropout=self.dropout, + ) + elif stack_type == "seasonality": + net_block = NBEATSSeasonalBlock( + units=self.widths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + min_period=self.expansion_coefficient_lengths[stack_id], + dropout=self.dropout, + ) + elif stack_type == "trend": + net_block = NBEATSTrendBlock( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + dropout=self.dropout, + ) + else: + raise ValueError(f"Unknown stack type {stack_type}") + + self.net_blocks.append(net_block) + + +NBeats_v2 = NBeats diff --git a/pytorch_forecasting/models/nbeats/_nbeatskan_pkg_v2.py b/pytorch_forecasting/models/nbeats/_nbeatskan_pkg_v2.py new file mode 100644 index 000000000..2fe8f5c50 --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeatskan_pkg_v2.py @@ -0,0 +1,81 @@ +"""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": ["jdb78", "Dev10-sys"], + "capability:exogenous": False, + "capability:multivariate": False, + "capability:pred_int": False, + "capability:flexible_history_length": False, + "capability:cold_start": False, + } + + @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 TslibDataModule + + return TslibDataModule + + @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]] = [ + { + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "backcast_loss_ratio": 1.0, + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + { + "stack_types": ["generic"], + "num_blocks": [1], + "num_block_layers": [2], + "widths": [16], + "expansion_coefficient_lengths": [8], + "sharing": [False], + }, + { + "loss": MAE(), + "widths": [16, 32], + "num_blocks": [1, 1], + "num_block_layers": [2, 2], + }, + ] + + default_dm_cfg: dict[str, Any] = { + "context_length": 8, + "prediction_length": 3, + "add_relative_time_idx": False, + } + + for param in params: + current_dm_cfg = param.get("datamodule_cfg", {}) + default_dm_cfg.update(current_dm_cfg) + param["datamodule_cfg"] = default_dm_cfg.copy() + + return params diff --git a/pytorch_forecasting/models/nbeats/_nbeatskan_v2.py b/pytorch_forecasting/models/nbeats/_nbeatskan_v2.py new file mode 100644 index 000000000..6aa55ce18 --- /dev/null +++ b/pytorch_forecasting/models/nbeats/_nbeatskan_v2.py @@ -0,0 +1,216 @@ +"""NBeatsKAN model for PyTorch Forecasting v2.""" + +from collections.abc import Callable +from typing import Any + +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.nbeats._nbeats_adapter_v2 import NBeatsAdapterV2 + + +class NBeatsKAN_v2(NBeatsAdapterV2): + """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. + backcast_loss_ratio : float, default=0.0 + Ratio of backcast loss to forecast loss. + 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, + backcast_loss_ratio: float = 0.0, + 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, + ): + 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"] + + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + metadata=metadata, + backcast_loss_ratio=backcast_loss_ratio, + ) + self.save_hyperparameters( + ignore=["loss", "logging_metrics", "metadata", "base_fun"] + ) + + self.stack_types = stack_types + self.num_blocks = num_blocks + self.num_block_layers = num_block_layers + self.widths = widths + self.sharing = sharing + self.expansion_coefficient_lengths = expansion_coefficient_lengths + self.dropout = dropout + + self.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._init_network() + + def _init_network(self): + """Build N-BEATS KAN stacks.""" + self.net_blocks = nn.ModuleList() + for stack_id, stack_type in enumerate(self.stack_types): + for _ in range(self.num_blocks[stack_id]): + net_block: nn.Module + if stack_type == "generic": + net_block = NBEATSGenericBlockKAN( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + dropout=self.dropout, + **self.kan_params, + ) + elif stack_type == "seasonality": + net_block = NBEATSSeasonalBlockKAN( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + nb_harmonics=None, # type: ignore[arg-type] + min_period=self.expansion_coefficient_lengths[stack_id], + dropout=self.dropout, + **self.kan_params, + ) + elif stack_type == "trend": + net_block = NBEATSTrendBlockKAN( + units=self.widths[stack_id], + thetas_dim=self.expansion_coefficient_lengths[stack_id], + num_block_layers=self.num_block_layers[stack_id], + backcast_length=self.context_length, + forecast_length=self.prediction_length, + dropout=self.dropout, + **self.kan_params, + ) + else: + raise ValueError(f"Unknown stack_type: {stack_type}") + + self.net_blocks.append(net_block)