diff --git a/tests/test_clustering.py b/tests/test_clustering.py index 8388af59..e09495f6 100644 --- a/tests/test_clustering.py +++ b/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) diff --git a/tslearn/clustering/kmeans.py b/tslearn/clustering/kmeans.py index e9097e4d..cef7eb4d 100644 --- a/tslearn/clustering/kmeans.py +++ b/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_ diff --git a/tslearn/clustering/utils.py b/tslearn/clustering/utils.py index a37e3d90..cf1b3ddc 100644 --- a/tslearn/clustering/utils.py +++ b/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,