tslearn_time_series_kmeans_clustering_with_soft_dtw.py

1import numpy
2from tslearn.datasets import CachedDatasets
3from tslearn.preprocessing import TimeSeriesScalerMeanVariance
4from tslearn.clustering import TimeSeriesKMeans
5
6# Seed for reproducibility
7seed = 0
8numpy.random.seed(seed)
9
10# Load a sample dataset
11X_train, y_train, X_test, y_test = CachedDatasets().load_dataset("Trace")
12
13# Keep only 50 time series for speed
14X_train = X_train[:50]
15numpy.random.shuffle(X_train)
16
17# Scale the time series to have zero mean and unit variance
18X_train = TimeSeriesScalerMeanVariance().fit_transform(X_train)
19
20# Soft-DTW-k-means
21print("Soft-DTW k-means")
22sdtw_km = TimeSeriesKMeans(n_clusters=3,
23                           metric="softdtw",
24                           metric_params={"gamma": .01},
25                           verbose=True,
26                           random_state=seed)
27y_pred = sdtw_km.fit_predict(X_train)
28
29print(y_pred)