ts2vec_classification_example.py

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

from benchpots.datasets import preprocess_random_walk
from pypots.classification import TS2Vec
from pypots.nn.functional import calc_binary_classification_metrics


def main():
    n_steps = 48
    n_features = 35
    n_classes = 2

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

    # 2. Extract training and test sets
    train_set = {"X": dataset["train_X"], "y": dataset["train_y"]}
    val_set = {"X": dataset["val_X"], "y": dataset["val_y"]}
    test_set = {"X": dataset["test_X"]}
    test_y_true = dataset["test_y"]

    # 3. Initialize the model
    model = TS2Vec(
        n_steps,
        n_features,
        n_classes=n_classes,
        n_output_dims=2,
        d_hidden=64,
        n_layers=2,
        epochs=5,
        device="cpu",
    )

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

    # 5. Predict classification labels
    print("🔮 Predicting classification labels...")
    results = model.predict(test_set)

    # 6. Evaluate accuracy
    proba_predictions = results["classification_proba"]
    metrics = calc_binary_classification_metrics(proba_predictions, test_y_true)
    print(
        f"✅ TS2Vec classification ROC_AUC: {metrics['roc_auc']:.4f}, "
        f"PR_AUC: {metrics['pr_auc']:.4f}, "
        f"F1: {metrics['f1']:.4f}, "
        f"Precision: {metrics['precision']:.4f}, "
        f"Recall: {metrics['recall']:.4f}"
    )


if __name__ == "__main__":
    main()