moment_imputation_example.py

"""
A minimalist, standalone example of the PyPOTS MOMENT model for time-series imputation.
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.imputation import MOMENT
from pypots.utils.metrics import calc_mse


def main():
    n_steps = 48
    n_features = 35

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

    # 2. Extract training and test sets
    train_set = {"X": dataset["train_X"], "X_ori": dataset["train_X_ori"]}
    val_set = {"X": dataset["val_X"], "X_ori": dataset["val_X_ori"]}
    test_set = {"X": dataset["test_X"], "X_ori": dataset["test_X_ori"]}
    test_X_intact = dataset["test_X_ori"]

    # 3. Initialize the model
    model = MOMENT(
        n_steps=n_steps,
        n_features=n_features,
        patch_size=2,
        patch_stride=2,
        n_layers=1,
        d_ffn=512,
        dropout=0.1,
        d_model=512,
        transformer_backbone="t5-small",
        transformer_type="encoder_only",
        head_dropout=0.1,
        finetuning_mode="zero-shot",
        revin_affine=False,
        add_positional_embedding=False,
        value_embedding_bias=False,
        orth_gain=1.41,
        mask_ratio=0,
        epochs=2,
        device="cpu",
    )

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

    # 5. Impute missing values
    print("🔮 Imputing missing values...")
    results = model.predict(test_set)
    imputed_X = results["imputation"]

    # 6. Evaluate
    indicating_mask = np.isnan(test_set["X"])
    mse = calc_mse(imputed_X, test_X_intact, indicating_mask)
    print(f"✅ The MSE of MOMENT imputation is: {mse:.4f}")


if __name__ == "__main__":
    main()