mixlinear_forecasting_example.py

"""
A minimalist, standalone example of the PyPOTS MixLinear model for time-series forecasting.
This script is auto-generated by extracting hyperparameters from the test code.
"""

import numpy as np
from benchpots.datasets import preprocess_random_walk
from pypots.nn.functional import calc_mse
from pypots.forecasting import MixLinear


def main():
    n_steps = 48
    n_pred_steps = 12
    n_features = 35

    # 1. Generate a random walk time-series dataset
    dataset = preprocess_random_walk(
        n_steps=n_steps + n_pred_steps, n_features=n_features, n_classes=5, n_samples_each_class=40, missing_rate=0.1
    )

    # 2. Extract training and test sets
    train_X = dataset["train_X"]
    val_X = dataset["val_X"]
    test_X = dataset["test_X"]
    train_set = {"X": train_X[:, :n_steps], "X_pred": train_X[:, n_steps:]}
    val_set = {"X": val_X[:, :n_steps], "X_pred": val_X[:, n_steps:]}
    test_set = {"X": test_X[:, :n_steps], "X_pred": test_X[:, n_steps:]}

    # 3. Initialize the model
    model = MixLinear(
        n_steps=n_steps,
        n_features=n_features,
        n_pred_steps=n_pred_steps,
        n_pred_features=n_features,
        period_len=2,
        lpf=1,
        alpha=0.5,
        rank=2,
        epochs=2,
        device="cpu",
    )

    # 4. Train the model
    print("🚀 Training the MixLinear forecasting model...")
    model.fit(train_set, val_set)

    # 5. Forecast
    print("🔮 Forecasting future steps...")
    results = model.predict(test_set)
    forecasts = results["forecasting"]

    test_MSE = calc_mse(forecasts, np.nan_to_num(test_set["X_pred"]), ~np.isnan(test_set["X_pred"]))
    print(f"✅ MixLinear forecasting MSE: {test_MSE:.4f}")


if __name__ == "__main__":
    main()