tslearn_timeseries_kmeans_clustering_with_soft_dtw.py

1import numpy as np
2from tslearn.clustering import TimeSeriesKMeans
3from tslearn.datasets import CachedDatasets
4from tslearn.preprocessing import TimeSeriesScalerMeanVariance
5
6# Load a sample dataset
7X_train, y_train, X_test, y_test = CachedDatasets().load_dataset("Trace")
8
9# Keep only 50 time series from the training set
10X_train = X_train[:50]
11
12# Normalize the time series
13X_train = TimeSeriesScalerMeanVariance().fit_transform(X_train)
14
15# Soft-DTW-k-means is a clustering method that relies on a differentiable
16# variant of DTW
17model = TimeSeriesKMeans(n_clusters=3, metric="softdtw", 
18                         metric_params={"gamma": .01},
19                         verbose=True, random_state=0)
20
21# Fit the model
22model.fit(X_train)
23
24# Predict clusters for the training data
25predictions = model.predict(X_train)
26
27print(f"Cluster assignments for the first 5 samples: {predictions[:5]}")