Fixes issue #268 TimeSeriesKmean with sample_weight

tests/test_clustering.py

@@ -236,6 +236,83 @@ def test_kmeans():
     np.testing.assert_equal(set(preds), set(range(4)))
 
 
+def test_kmeans_sample_weight():
+    """Test sample weights for Euclidean TimeSeriesKMeans."""
+    X = to_time_series_dataset([
+        [0.0, 0.0],
+        [10.0, 10.0],
+        [20.0, 20.0],
+    ])
+
+    sample_weight = np.array([1.0, 1.0, 10.0])
+
+    km = TimeSeriesKMeans(
+        n_clusters=1,
+        metric="euclidean",
+        max_iter=5,
+        random_state=0,
+    ).fit(X, sample_weight=sample_weight)
+
+    expected_center = np.average(
+        X,
+        axis=0,
+        weights=sample_weight,
+    )
+
+    np.testing.assert_allclose(
+        km.cluster_centers_[0],
+        expected_center,
+    )
+
+    # Check weighted inertia
+    distances = cdist(
+        X.reshape((3, -1)),
+        km.cluster_centers_.reshape((1, -1)),
+    ).ravel()
+
+    expected_inertia = np.average(
+        distances ** 2,
+        weights=sample_weight,
+    )
+
+    np.testing.assert_allclose(
+        km.inertia_,
+        expected_inertia,
+    )
+
+def test_kmeans_sample_weight_uniform():
+    """Uniform sample weights should give the same result as no weights."""
+    rng = np.random.RandomState(0)
+    X = rng.randn(15, 10, 3)
+
+    km = TimeSeriesKMeans(
+        n_clusters=3,
+        metric="euclidean",
+        max_iter=5,
+        random_state=0,
+    ).fit(X)
+
+    km_weighted = TimeSeriesKMeans(
+        n_clusters=3,
+        metric="euclidean",
+        max_iter=5,
+        random_state=0,
+    ).fit(X, sample_weight=np.ones(X.shape[0]))
+
+    np.testing.assert_allclose(
+        km.cluster_centers_,
+        km_weighted.cluster_centers_,
+    )
+    np.testing.assert_equal(
+        km.labels_,
+        km_weighted.labels_,
+    )
+    np.testing.assert_allclose(
+        km.inertia_,
+        km_weighted.inertia_,
+    )
+
+
 def test_kshape():
     n, sz, d = 15, 10, 3
     rng = np.random.RandomState(0)

tslearn/clustering/kmeans.py

@@ -639,14 +639,15 @@ def _get_metric_params(self):
             del metric_params["n_jobs"]
         return metric_params
 
-    def _fit_one_init(self, X, x_squared_norms, rs):
+    def _fit_one_init(self, X, x_squared_norms, rs, sample_weight):
         metric_params = self._get_metric_params()
         n_ts, sz, d = X.shape
         if hasattr(self.init, "__array__"):
             self.cluster_centers_ = self.init.copy()
         elif isinstance(self.init, str) and self.init == "k-means++":
             if self.metric == "euclidean":
-                sample_weight = _check_sample_weight(None, X, dtype=X.dtype)
+                # sample_weight = _check_sample_weight(None, X, dtype=X.dtype)
+                sample_weight = self.sample_weight_
                 self.cluster_centers_ = _kmeans_plusplus(
                     X.reshape((n_ts, -1)),
                     self.n_clusters,
@@ -697,7 +698,8 @@ def metric_fun(x, y):
             self._assign(X)
             if self.verbose:
                 print("%.3f" % self.inertia_, end=" --> ")
-            self._update_centroids(X)
+
+            self._update_centroids(X, sample_weight)
 
             if numpy.abs(old_inertia - self.inertia_) < self.tol:
                 break
@@ -752,36 +754,49 @@ def _assign(self, X, update_class_attributes=True):
             else:
                 inertia_dists = dists
             self.inertia_ = _compute_inertia(
-                inertia_dists, self.labels_, self._squared_inertia
+                inertia_dists,
+                self.labels_,
+                self._squared_inertia,
+                sample_weight=self.sample_weight_,
             )
         return matched_labels
 
-    def _update_centroids(self, X):
+    def _update_centroids(self, X,sample_weight):
         metric_params = self._get_metric_params()
         for k in range(self.n_clusters):
+            mask = self.labels_ == k
+            X_k = X[mask]
+            weights_k = sample_weight[mask]
+
             if self.metric == "dtw":
                 self.cluster_centers_[k] = dtw_barycenter_averaging_petitjean(
-                    X=X[self.labels_ == k],
+                    X=X_k,
                     barycenter_size=None,
                     init_barycenter=self.cluster_centers_[k],
                     metric_params=metric_params,
                     verbose=False,
-                    n_jobs=self.n_jobs
+                    n_jobs=self.n_jobs,
+                    weights=weights_k,
                 )
+
             elif self.metric == "softdtw":
                 self.cluster_centers_[k] = softdtw_barycenter(
-                    X=X[self.labels_ == k],
+                    X=X_k,
                     max_iter=self.max_iter_barycenter,
                     init=self.cluster_centers_[k],
+                    weights=weights_k,
                     n_jobs=self.n_jobs,
                     **metric_params
                 )
+
             else:
-                # Euclidean
-                self.cluster_centers_[k] = numpy.average(X[self.labels_ == k],
-                                                         axis=0)
+                self.cluster_centers_[k] = numpy.average(
+                    X_k,
+                    axis=0,
+                    weights=weights_k,
+                )
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None, sample_weight=None):
         """Compute k-means clustering.
 
         Parameters
@@ -791,6 +806,10 @@ def fit(self, X, y=None):
 
         y
             Ignored
+
+        sample_weight : array-like of shape=(n_ts, ) or None (default: None)
+            Weights to be given to time series in the learning process. By
+            default, all time series weights are equal.
         """
 
         X = check_array(
@@ -819,6 +838,8 @@ def fit(self, X, y=None):
         max_attempts = max(self.n_init, 10)
 
         X_ = to_time_series_dataset(X)
+        sample_weight = _check_sample_weight(sample_weight, X_)
+        self.sample_weight_ = sample_weight
         rs = check_random_state(self.random_state)
 
         if (
@@ -843,7 +864,7 @@ def fit(self, X, y=None):
                 if self.verbose and self.n_init > 1:
                     print("Init %d" % (n_successful + 1))
                 n_attempts += 1
-                self._fit_one_init(X_, x_squared_norms, rs)
+                self._fit_one_init(X_, x_squared_norms, rs, sample_weight)
                 if self.inertia_ < min_inertia:
                     best_correct_centroids = self.cluster_centers_.copy()
                     min_inertia = self.inertia_

tslearn/clustering/utils.py

@@ -48,23 +48,42 @@ def _check_full_length(centroids):
     return resampler.fit_transform(centroids)
 
 
-def _compute_inertia(distances, assignments, squared=True):
+def _compute_inertia(
+    distances, assignments, squared=True, sample_weight=None
+):
     """Derive inertia (average of squared distances) from pre-computed
     distances and assignments.
 
-    Examples
-    --------
-    >>> dists = numpy.array([[1., 2., 0.5], [0., 3., 1.]])
-    >>> assign = numpy.array([2, 0])
-    >>> float(_compute_inertia(dists, assign))
-    0.125
+    Parameters
+    ----------
+    distances : array-like of shape=(n_ts, n_clusters)
+        Pre-computed distances.
+
+    assignments : array-like of shape=(n_ts,)
+        Cluster assignment for each time series.
+
+    squared : bool (default=True)
+        Whether to square the distances before averaging.
+
+    sample_weight : array-like of shape=(n_ts,) or None (default=None)
+        Weight of each time series.
+
+    Returns
+    -------
+    inertia : float
+        Weighted average of distances (or squared distances).
+
     """
     n_ts = distances.shape[0]
+    assigned_distances = distances[numpy.arange(n_ts), assignments]
+
     if squared:
-        return numpy.sum(distances[numpy.arange(n_ts),
-                                   assignments] ** 2) / n_ts
-    else:
-        return numpy.sum(distances[numpy.arange(n_ts), assignments]) / n_ts
+        assigned_distances = assigned_distances ** 2
+
+    if sample_weight is None:
+        return numpy.sum(assigned_distances) / n_ts
+
+    return numpy.average(assigned_distances, weights=sample_weight)
 
 
 def silhouette_score(X, labels, metric=None, sample_size=None,