def test_init_deterministic(self, init, dtype_device):
dtype, device = dtype_device
random_state = 2
x = torch.rand((20, 5), dtype=dtype, device=device)
kmeans_layer = KMeans(n_clusters=3, init=init)
torch.manual_seed(random_state)
cluster_centers_1 = kmeans_layer.initialize(x)
torch.manual_seed(random_state)
cluster_centers_2 = kmeans_layer.initialize(x)
assert torch.allclose(cluster_centers_1, cluster_centers_2)
def test_dummy(self, dtype_device, random_state):
"""Create a dummy feature matrix.
Notes
-----
Copied from scikit-learn tests.
"""
dtype, device = dtype_device
x = torch.tensor([[0, 0], [0.5, 0], [0.5, 1], [1,
1]]).to(dtype=dtype,
device=device)
manual_init = torch.tensor([[0, 0], [1, 1]])
kmeans_layer = KMeans(n_clusters=2,
init='manual',
random_state=random_state,
n_init=1,
verbose=True)
cluster_ixs, cluster_centers = kmeans_layer(x, manual_init=manual_init)
assert torch.allclose(cluster_ixs,
torch.tensor([0, 0, 1, 1]).to(device=device))
assert torch.allclose(
cluster_centers,
torch.tensor([[0.25, 0], [0.75, 1]]).to(dtype=dtype,
device=device))
def test_compute_distances(self, dtype_device):
dtype, device = dtype_device
x = torch.tensor([[1, 0], [0, 1], [2, 2]]).to(dtype=dtype, device=device)
cluster_centers = torch.tensor([[0, 0], [1, 1]]).to(dtype=dtype, device=device)
correct_result = torch.tensor([[1, 1], [1, 1], [8, 2]]).to(dtype=dtype, device=device)
assert torch.allclose(KMeans.compute_distances(x, cluster_centers), correct_result)
def test_all_deterministic(self, init, dtype_device):
dtype, device = dtype_device
random_state = 2
kmeans_layer = KMeans(n_clusters=3, init=init, random_state=random_state, n_init=2, verbose=True)
x = torch.rand((20, 5), dtype=dtype, device=device)
cluster_ixs_1, cluster_centers_1 = kmeans_layer(x)
cluster_ixs_2, cluster_centers_2 = kmeans_layer(x)
assert torch.allclose(cluster_centers_1, cluster_centers_2)
assert torch.allclose(cluster_ixs_1, cluster_ixs_2)
def test_manual_init(self):
n_samples = 10
n_features = 3
n_clusters = 2
kmeans_layer = KMeans(init='manual', n_clusters=n_clusters)
x = torch.rand((n_samples, n_features))
manual_init = torch.rand((n_clusters, n_features))
wrong_init_1 = torch.rand((n_clusters + 1, n_features))
wrong_init_2 = torch.rand((n_clusters, n_features + 1))
with pytest.raises(TypeError):
kmeans_layer.initialize(x, manual_init=None)
with pytest.raises(ValueError):
kmeans_layer.initialize(x, manual_init=wrong_init_1)
with pytest.raises(ValueError):
kmeans_layer.initialize(x, manual_init=wrong_init_2)
assert torch.allclose(manual_init, kmeans_layer.initialize(x, manual_init=manual_init))
def test_errors(self):
with pytest.raises(ValueError):
KMeans(init='fake')
with pytest.raises(ValueError):
KMeans(n_clusters=4)(torch.ones(3, 2))