diff --git a/heat/cluster/batchparallelclustering.py b/heat/cluster/batchparallelclustering.py index de795cdb89..d6d64e3b1c 100644 --- a/heat/cluster/batchparallelclustering.py +++ b/heat/cluster/batchparallelclustering.py @@ -42,7 +42,25 @@ def _initialize_plus_plus( for i in range(1, n_clusters): dist = torch.cdist(X, X[idxs[:i]], p=p) dist = torch.min(dist, dim=1)[0] - idxs[i] = torch.multinomial(weights * dist, 1) + probs = weights * dist + probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0) + + # Minimal fallback ONLY if multinomial would crash + if probs.sum() <= 0: + # fall back to standard k-means++ (ignore weights) + probs = torch.nan_to_num(dist, nan=0.0, posinf=0.0, neginf=0.0) + + if probs.sum() <= 0: + # fully degenerate (all distances zero) -> pick any not-yet-picked index if possible + mask = torch.ones(X.shape[0], dtype=torch.bool, device=X.device) + mask[idxs[:i]] = False + candidates = torch.nonzero(mask, as_tuple=False).flatten() + if candidates.numel() > 0: + idxs[i] = candidates[torch.randint(0, candidates.numel(), (1,), device=X.device)] + else: + idxs[i] = torch.randint(0, X.shape[0], (1,), device=X.device) + else: + idxs[i] = torch.multinomial(probs, 1) return X[idxs] diff --git a/tests/cluster/test_batchparallelclustering.py b/tests/cluster/test_batchparallelclustering.py index 53bb50a40e..99c3c807ae 100644 --- a/tests/cluster/test_batchparallelclustering.py +++ b/tests/cluster/test_batchparallelclustering.py @@ -39,8 +39,42 @@ def test_kmex(self): _kmex(X, 2, 2, init, max_iter, tol) def test_initialize_plus_plus(self): - X = torch.rand(100, 3) - _initialize_plus_plus(X, 3, 2, random_state=None, max_samples=50) + with self.subTest("subsampling"): + X = torch.rand(100, 3) + centers = _initialize_plus_plus(X, 3, 2, random_state=0, max_samples=50) + self.assertEqual(centers.shape, (3, 3)) + + # 2) probs.sum() <= 0 because weights are all zero -> fallback to dist -> multinomial runs + with self.subTest("weights_zero_fallback_to_dist"): + X = torch.rand(30, 3) + weights = torch.zeros(X.shape[0], dtype=X.dtype) + centers = _initialize_plus_plus(X, 3, 2, random_state=0, weights=weights) + self.assertEqual(centers.shape, (3, 3)) + + # 3) fully degenerate distances (all points identical) -> probs.sum() <= 0 twice -> candidate selection branch + with self.subTest("all_distances_zero_candidate_selection"): + X = torch.ones(10, 3) + weights = torch.ones(X.shape[0], dtype=X.dtype) + centers = _initialize_plus_plus(X, 3, 2, random_state=0, weights=weights) + self.assertEqual(centers.shape, (3, 3)) + + # 4) extreme degenerate case: only one sample, n_clusters>1 -> candidates empty branch + with self.subTest("single_sample_candidates_empty"): + X = torch.ones(1, 3) + centers = _initialize_plus_plus(X, 2, 2, random_state=0) + self.assertEqual(centers.shape, (2, 3)) + + # 5) NaN-handling path -> nan_to_num is exercised (should not crash) + with self.subTest("nan_to_num_path"): + X = torch.tensor( + [[0.0, 0.0, 0.0], + [float("nan"), 0.0, 0.0], + [1.0, 0.0, 0.0]], + dtype=torch.float32, + ) + # seed chosen so first centroid is deterministic (helps avoid flakiness) + centers = _initialize_plus_plus(X, 2, 2, random_state=2) + self.assertEqual(centers.shape, (2, 3)) def test_BatchParallelKClustering(self): with self.assertRaises(TypeError):