diff --git a/docs/source/m_layer_v2.rst b/docs/source/m_layer_v2.rst index 80b447f95..9dae3da33 100644 --- a/docs/source/m_layer_v2.rst +++ b/docs/source/m_layer_v2.rst @@ -47,6 +47,7 @@ See the detailed API documentation for the V2 base classes and specific model im models.samformer._samformer_v2.Samformer models.tide._tide_dsipts._tide_v2.TIDE models.timexer._timexer_v2.TimeXer + models.tsmixer._tsmixer_v2.TSMixer models.xlstm._xlstm_v2.xLSTMTime_v2 models.mlp._decodermlp_v2.DecoderMLP_v2 models.softs._softs_v2.SOFTS diff --git a/docs/source/pkg_v2.rst b/docs/source/pkg_v2.rst index ef296de6b..24504ea62 100644 --- a/docs/source/pkg_v2.rst +++ b/docs/source/pkg_v2.rst @@ -99,6 +99,7 @@ See the detailed API documentation for the available V2 Package classes below: models.samformer._samformer_v2_pkg.Samformer_pkg_v2 models.tide._tide_dsipts._tide_v2_pkg.TIDE_pkg_v2 models.timexer._timexer_pkg_v2.TimeXer_pkg_v2 + models.tsmixer._tsmixer_pkg_v2.TSMixer_pkg_v2 models.xlstm._xlstm_pkg_v2.xLSTMTime_pkg_v2 models.mlp._decodermlp_pkg_v2.DecoderMLP_pkg_v2 models.softs._softs_pkg_v2.SOFTS_pkg_v2 diff --git a/pytorch_forecasting/models/tsmixer/__init__.py b/pytorch_forecasting/models/tsmixer/__init__.py new file mode 100644 index 000000000..5d2fbdc26 --- /dev/null +++ b/pytorch_forecasting/models/tsmixer/__init__.py @@ -0,0 +1,8 @@ +""" +TSMixer model for time series forecasting. +""" + +from pytorch_forecasting.models.tsmixer._tsmixer_pkg_v2 import TSMixer_pkg_v2 +from pytorch_forecasting.models.tsmixer._tsmixer_v2 import TSMixer + +__all__ = ["TSMixer", "TSMixer_pkg_v2"] diff --git a/pytorch_forecasting/models/tsmixer/_tsmixer_pkg_v2.py b/pytorch_forecasting/models/tsmixer/_tsmixer_pkg_v2.py new file mode 100644 index 000000000..b6a812715 --- /dev/null +++ b/pytorch_forecasting/models/tsmixer/_tsmixer_pkg_v2.py @@ -0,0 +1,85 @@ +""" +Packages container for TSMixer model. +""" + +from pytorch_forecasting.base._base_pkg import Base_pkg + + +class TSMixer_pkg_v2(Base_pkg): + """TSMixer package container.""" + + _tags = { + "info:name": "TSMixer", + "info:compute": 2, + "authors": ["seaic-mac-murchadha"], + "info:y_type": ["numeric"], + "capability:exogenous": True, + "capability:multivariate": True, + "capability:pred_int": True, + "capability:flexible_history_length": False, + "capability:cold_start": False, + } + + @classmethod + def get_cls(cls): + """Get model class.""" + from pytorch_forecasting.models.tsmixer._tsmixer_v2 import TSMixer + + return TSMixer + + @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[dict] + Parameter configurations used to create testing instances of the TSMixer + class. + """ + + import torch.nn as nn + + from pytorch_forecasting.metrics import SMAPE + + params = [ + {}, + dict( + d_model=64, + e_layers=2, + dropout=0.1, + logging_metrics=[SMAPE()], + ), + dict( + d_model=32, + e_layers=1, + dropout=0.0, + loss=nn.MSELoss(), + ), + dict( + optimizer="adamw", + optimizer_params={"lr": 1e-3}, + ), + ] + + 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 diff --git a/pytorch_forecasting/models/tsmixer/_tsmixer_v2.py b/pytorch_forecasting/models/tsmixer/_tsmixer_v2.py new file mode 100644 index 000000000..81a2a8306 --- /dev/null +++ b/pytorch_forecasting/models/tsmixer/_tsmixer_v2.py @@ -0,0 +1,279 @@ +""" +TSMixer model for PyTorch Forecasting. +------------------------------------------- +""" + +################################################# +# NOTE: This is an experimental implementation # +# of TSMixer for PyTorch Forecasting v2. # +# It is an unstable API and subject to change. # +################################################# + +from typing import Any + +import torch +import torch.nn as nn +from torch.optim import Optimizer + +from pytorch_forecasting.metrics import QuantileLoss +from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel + + +class TSMixerBlock(nn.Module): + """ + TSMixer block for applying the time-mixing and feature-mixing MLPs. + + Parameters + ---------- + sequence_length : int + Length of the lookback window containing past time steps. + num_features : int + Number of expected features in the input. + hidden_dim : int + Dimension of the hidden layers. + dropout : float + Probability of an element to be zeroed. + """ + + def __init__( + self, + sequence_length: int, + num_features: int, + hidden_dim: int, + dropout: float, + ): + super().__init__() + + self.temporal = nn.Sequential( + nn.Linear(sequence_length, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, sequence_length), + nn.Dropout(dropout), + ) + + self.channel = nn.Sequential( + nn.Linear(num_features, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, num_features), + nn.Dropout(dropout), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Apply time-mixing and feature-mixing to the input tensor. + + Parameters + ---------- + x : torch.Tensor + Input tensor. + + Returns + ------- + torch.Tensor + Output tensor with time-mixing and feature-mixing applied. + """ + x = x + self.temporal(x.transpose(1, 2)).transpose(1, 2) + x = x + self.channel(x) + + return x + + +class TSMixer(TslibBaseModel): + """ + TSMixer: An All-MLP Architecture for Time Series Forecasting. + + TSMixer is a lightweight time series forecasting model made of stacked + multilayer perceptrons. + + Parameters + ---------- + loss : nn.Module + Loss function for training step. + d_model : int, default=64 + Hidden dimension of the temporal and channel mixing MLPs. + e_layers : int, default=2 + Number of stacked TSMixer blocks. + dropout : float, default=0.1 + Dropout probability used within the mixer blocks. + logging_metrics : Optional[list[nn.Module]], default=None + List of metrics to log during training, validation and testing. + optimizer : Optional[Union[Optimizer, str]], default="adam" + Optimizer used for training. + optimizer_params : Optional[dict], default=None + Parameters passed to the optimizer. + lr_scheduler : Optional[str], default=None + Learning rate scheduler utilised. + lr_scheduler_params : Optional[dict], default=None + Parameters passed to the learning rate scheduler. + metadata : Optional[dict], default=None + Metadata for the model from TslibDataModule. + + References + ---------- + [1] TSMixer: An All-MLP Architecture for Time Series Forecasting (https://arxiv.org/abs/2303.06053). + [2] https://github.com/thuml/Time-Series-Library/blob/main/models/TSMixer.py + """ + + @classmethod + def _pkg(cls): + """Package containing the model.""" + from pytorch_forecasting.models.tsmixer._tsmixer_pkg_v2 import TSMixer_pkg_v2 + + return TSMixer_pkg_v2 + + def __init__( + self, + loss: nn.Module, + d_model: int = 64, + e_layers: int = 2, + 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, + metadata: dict | None = None, + **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.d_model = d_model + self.e_layers = e_layers + self.dropout = dropout + + self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"]) + + self._init_network() + + def _init_network(self): + """Initialize the TSMixer network components.""" + + self.enc_in = self.cont_dim + self.target_dim + + self.model = nn.ModuleList( + [ + TSMixerBlock( + sequence_length=self.context_length, + num_features=self.enc_in, + hidden_dim=self.d_model, + dropout=self.dropout, + ) + for _ in range(self.e_layers) + ] + ) + + self.n_quantiles = None + output_dim = self.prediction_length + + if isinstance(self._loss, QuantileLoss): + self.n_quantiles = len(self._loss.quantiles) + output_dim *= self.n_quantiles + + self.projection = nn.Linear( + self.context_length, + output_dim, + ) + + def _encoder( + self, x: torch.Tensor, target_indices: torch.Tensor | None + ) -> torch.Tensor: + """ + Encode the input time series through the TSMixer blocks. + + Parameters + ---------- + x : torch.Tensor + Input data to the encoder via tensor of shape + (batch_size, context_length, n_features). + target_indices : torch.Tensor or None + Indices of target features to extract from the output. + If None, all features are returned. + + Returns + ------- + torch.Tensor + Forecast tensor. + For point forecasts, returns shape + (batch_size, prediction_length, n_targets). + For quantile forecasts, returns shape + (batch_size, prediction_length, n_quantiles). + """ + + for block in self.model: + x = block(x) + + output = self.projection(x.transpose(1, 2)) + + if target_indices is not None: + output = output[:, target_indices, :] + + if self.n_quantiles is None: + return output.transpose(1, 2) + + if output.shape[1] != 1: + raise ValueError( + "Quantile forecasting currently only supports a single target." + ) + + return output.squeeze(1).reshape( + output.shape[0], + self.prediction_length, + self.n_quantiles, + ) + + def _prepare_input_data(self, x: dict[str, torch.Tensor]): + """Prepare input data and target indices for model input.""" + + available_features = [] + current_idx = 0 + + if "history_cont" in x and x["history_cont"].size(-1) > 0: + available_features.append(x["history_cont"]) + current_idx += x["history_cont"].size(-1) + + if "history_target" not in x or x["history_target"].size(-1) == 0: + raise ValueError("No target history found in the input dictionary.") + + n_targets = x["history_target"].size(-1) + target_indices = list(range(current_idx, current_idx + n_targets)) + available_features.append(x["history_target"]) + + input_data = torch.cat(available_features, dim=-1) + + target_indices = torch.tensor( + target_indices, dtype=torch.long, device=input_data.device + ) + + return input_data, target_indices + + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Forward pass of the TSMixer model. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Dictionary containing input tensors. + + Returns + ------- + dict[str, torch.Tensor] + Dictionary containing the model prediction. + """ + + input_data, target_indices = self._prepare_input_data(x) + + prediction = self._encoder(input_data, target_indices) + + if "target_scale" in x and hasattr(self, "transform_output"): + prediction = self.transform_output(prediction, x["target_scale"]) + + return {"prediction": prediction} diff --git a/tests/test_models/test_tsmixer_v2.py b/tests/test_models/test_tsmixer_v2.py new file mode 100644 index 000000000..b23ffd228 --- /dev/null +++ b/tests/test_models/test_tsmixer_v2.py @@ -0,0 +1,233 @@ +import numpy as np +import pandas as pd +import pytest +import torch + +from pytorch_forecasting.data import TimeSeries +from pytorch_forecasting.data.data_module import TslibDataModule +from pytorch_forecasting.metrics import MAE, SMAPE, QuantileLoss +from pytorch_forecasting.models.tsmixer._tsmixer_v2 import TSMixer + + +@pytest.fixture +def sample_dataset(): + """Create a sample dataset for testing using v2.""" + n_samples = 100 + n_series = 3 + + time_idx = np.arange(n_samples) + + series_data = [] + for i in range(n_series): + trend = 0.1 * time_idx + seasonality = 10 * np.sin(2 * np.pi * time_idx / 20) + noise = np.random.normal(0, 1, n_samples) + values = trend + seasonality + noise + + series = pd.DataFrame( + { + "time_idx": time_idx, + "series_id": i, + "value": values, + "feat1": np.random.normal(0, 1, n_samples), + "feat2": np.random.normal(0, 1, n_samples), + } + ) + series_data.append(series) + + data = pd.concat(series_data).reset_index(drop=True) + + ts = TimeSeries( + data, + time="time_idx", + group=["series_id"], + target=["value"], + num=["feat1", "feat2"], + cat=[], + known=["time_idx"], + unknown=["value", "feat1", "feat2"], + ) + + dm = TslibDataModule(ts, context_length=16, prediction_length=4, batch_size=4) + + dm.setup() + + return {"data_module": dm, "time_series": ts} + + +def test_quantile_loss_output(sample_dataset): + """Test TSMixer output shape with quantile loss.""" + + dm = sample_dataset["data_module"] + + train_dataloader = dm.train_dataloader() + batch = next(iter(train_dataloader))[0] + + metadata = dm.metadata + + quantiles = [0.1, 0.5, 0.9] + + model = TSMixer( + loss=QuantileLoss(quantiles=quantiles), + d_model=32, + e_layers=2, + dropout=0.1, + logging_metrics=[SMAPE(), MAE()], + metadata=metadata, + ) + + with torch.no_grad(): + output = model(batch) + + assert "prediction" in output + pred = output["prediction"] + assert pred.shape == ( + dm.batch_size, + metadata["prediction_length"], + len(quantiles), + ) + + +def test_univariate_forecast(): + """Test univariate forecasting with TSMixer.""" + + n_samples = 100 + time_idx = np.arange(n_samples) + values = np.sin(2 * np.pi * time_idx / 20) + np.random.normal(0, 0.1, n_samples) + + series = pd.DataFrame({"time_idx": time_idx, "series_id": 0, "value": values}) + + ts = TimeSeries( + series, + time="time_idx", + group=["series_id"], + target=["value"], + num=[], + cat=[], + known=["time_idx"], + unknown=["value"], + ) + + dm = TslibDataModule(ts, context_length=16, prediction_length=4, batch_size=4) + + dm.setup() + + metadata = dm.metadata + + model = TSMixer( + loss=MAE(), + d_model=32, + e_layers=1, + dropout=0.1, + metadata=metadata, + ) + + train_dataloader = dm.train_dataloader() + batch = next(iter(train_dataloader))[0] + + with torch.no_grad(): + output = model(batch) + + assert "prediction" in output + assert output["prediction"].shape == ( + dm.batch_size, + metadata["prediction_length"], + 1, + ) + + +def test_prepare_input_data(sample_dataset): + """Test preparation of continuous and target historical data.""" + + dm = sample_dataset["data_module"] + batch = next(iter(dm.train_dataloader()))[0] # One sample batch + + model = TSMixer( + loss=MAE(), + metadata=dm.metadata, + ) + + input_data, target_indices = model._prepare_input_data(batch) + + assert input_data.shape[-1] == ( + batch["history_cont"].shape[-1] + batch["history_target"].shape[-1] + ) + + assert target_indices.tolist() == [batch["history_cont"].shape[-1]] + + assert target_indices.dtype == torch.long + assert target_indices.device == input_data.device + + +def test_prepare_input_data_error_for_no_history(sample_dataset): + """Test _prepare_input_data without the required history.""" + + dm = sample_dataset["data_module"] + batch = next(iter(dm.train_dataloader()))[0] + + model = TSMixer( + loss=MAE(), + metadata=dm.metadata, + ) + + batch_without_target = { + key: value for key, value in batch.items() if key != "history_target" + } + + with pytest.raises( + ValueError, + match="No target history found in the input dictionary.", + ): + model._prepare_input_data(batch_without_target) + + +def test_quantile_loss_error_on_multiple_targets(sample_dataset): + """Test error for quantile forecasting with multiple targets.""" + + dm = sample_dataset["data_module"] + + model = TSMixer( + loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]), + metadata=dm.metadata, + ) + + batch_size = dm.batch_size + context_length = dm.metadata["context_length"] + n_features = model.enc_in + + x = torch.randn(batch_size, context_length, n_features) + + target_indices = torch.tensor([0, 1], dtype=torch.long) + + with pytest.raises( + ValueError, + match="Quantile forecasting currently only supports a single target.", + ): + model._encoder(x, target_indices) + + +def test_forward_with_target_scale(sample_dataset): + """Test that prediction uses target scaling when target_scale is provided.""" + dm = sample_dataset["data_module"] + batch = next(iter(dm.train_dataloader()))[0] + + model = TSMixer( + loss=MAE(), + metadata=dm.metadata, + ) + + batch["target_scale"] = torch.ones(batch["history_target"].shape[0], 2) + + captured = {} + + def mock_transform_output(prediction, target_scale): + captured["prediction"] = prediction + captured["target_scale"] = target_scale + return prediction + 1 + + model.transform_output = mock_transform_output + + prediction = model(batch)["prediction"] + + assert torch.equal(captured["target_scale"], batch["target_scale"]) + assert torch.allclose(prediction, captured["prediction"] + 1)