Fix NMI/AMI returning 0.0 instead of 1.0 for two single-cluster labelings

CHANGELOG.md

@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed `normalized_mutual_info_score` and `adjusted_mutual_info_score` returning `0.0` instead of `1.0` for two single-cluster labelings ([#3470](https://github.com/Lightning-AI/torchmetrics/pull/3470))
 - Fixed malformed LaTeX in `CLIPScore` and `HausdorffDistance` docstring math so it renders correctly ([#3427](https://github.com/Lightning-AI/torchmetrics/pull/3427))
 
 

src/torchmetrics/functional/clustering/adjusted_mutual_info_score.py

@@ -21,6 +21,7 @@
     _validate_average_method_arg,
     calculate_entropy,
     calculate_generalized_mean,
+    check_cluster_labels,
 )
 
 
@@ -46,6 +47,14 @@ def adjusted_mutual_info_score(
 
     """
     _validate_average_method_arg(average_method)
+    check_cluster_labels(preds, target)
+
+    # Special limit case: both labelings have a single cluster (zero entropy). Since the two
+    # labelings then trivially agree, this is a perfect match and the score is 1.0. Matches the
+    # convention used by :func:`sklearn.metrics.adjusted_mutual_info_score`.
+    if preds.unique().numel() == target.unique().numel() == 1:
+        return torch.tensor(1.0, dtype=torch.float32, device=preds.device)
+
     contingency = _mutual_info_score_update(preds, target)
     mutual_info = _mutual_info_score_compute(contingency)
     expected_mutual_info = expected_mutual_info_score(contingency, target.numel())

src/torchmetrics/functional/clustering/normalized_mutual_info_score.py

@@ -48,6 +48,13 @@ def normalized_mutual_info_score(
     """
     check_cluster_labels(preds, target)
     _validate_average_method_arg(average_method)
+
+    # Special limit case: both labelings have a single cluster (zero entropy). Since the two
+    # labelings then trivially agree, this is a perfect match and the score is 1.0. Matches the
+    # convention used by :func:`sklearn.metrics.normalized_mutual_info_score`.
+    if preds.unique().numel() == target.unique().numel() == 1:
+        return torch.tensor(1.0, dtype=torch.float32, device=preds.device)
+
     mutual_info = mutual_info_score(preds, target)
     if torch.allclose(mutual_info, torch.tensor(0.0), atol=torch.finfo().eps):
         return mutual_info

tests/unittests/clustering/test_adjusted_mutual_info_score.py

@@ -77,6 +77,19 @@ def test_adjusted_mutual_info_score_functional_single_cluster(average_method):
     assert torch.allclose(adjusted_mutual_info_score(tensor_b, tensor_a, average_method), torch.tensor(0.0), atol=ATOL)
 
 
+@pytest.mark.parametrize("average_method", ["min", "geometric", "arithmetic", "max"])
+def test_adjusted_mutual_info_score_functional_both_single_cluster(average_method):
+    """Check that two identical single-cluster labelings give a perfect score of 1, matching sklearn."""
+    tensor_a = torch.zeros((BATCH_SIZE,), dtype=torch.int)
+    tensor_b = torch.zeros((BATCH_SIZE,), dtype=torch.int)
+    expected = sklearn_ami(tensor_a.numpy(), tensor_b.numpy(), average_method=average_method)
+    assert torch.allclose(
+        adjusted_mutual_info_score(tensor_a, tensor_b, average_method),
+        torch.tensor(expected, dtype=torch.float32),
+        atol=ATOL,
+    )
+
+
 @pytest.mark.parametrize("average_method", ["min", "geometric", "arithmetic", "max"])
 def test_adjusted_mutual_info_score_functional_raises_invalid_task(average_method):
     """Check that metric rejects continuous-valued inputs."""

tests/unittests/clustering/test_normalized_mutual_info_score.py

@@ -75,6 +75,17 @@ def test_normalized_mutual_info_score_functional_single_cluster(average_method):
     assert torch.allclose(normalized_mutual_info_score(tensor_b, tensor_a, average_method), torch.tensor(0.0))
 
 
+@pytest.mark.parametrize("average_method", ["min", "geometric", "arithmetic", "max"])
+def test_normalized_mutual_info_score_functional_both_single_cluster(average_method):
+    """Check that two identical single-cluster labelings give a perfect score of 1, matching sklearn."""
+    tensor_a = torch.zeros((BATCH_SIZE,), dtype=torch.int)
+    tensor_b = torch.zeros((BATCH_SIZE,), dtype=torch.int)
+    expected = sklearn_nmi(tensor_a.numpy(), tensor_b.numpy(), average_method=average_method)
+    assert torch.allclose(
+        normalized_mutual_info_score(tensor_a, tensor_b, average_method), torch.tensor(expected, dtype=torch.float32)
+    )
+
+
 @pytest.mark.parametrize("average_method", ["min", "geometric", "arithmetic", "max"])
 def test_normalized_mutual_info_score_functional_raises_invalid_task(average_method):
     """Check that metric rejects continuous-valued inputs."""