diff --git a/CHANGELOG.md b/CHANGELOG.md index 9518750113a..203485870a8 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/src/torchmetrics/functional/clustering/adjusted_mutual_info_score.py b/src/torchmetrics/functional/clustering/adjusted_mutual_info_score.py index b70525c1d76..8665a17375a 100644 --- a/src/torchmetrics/functional/clustering/adjusted_mutual_info_score.py +++ b/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()) diff --git a/src/torchmetrics/functional/clustering/normalized_mutual_info_score.py b/src/torchmetrics/functional/clustering/normalized_mutual_info_score.py index 5cdc81b960d..2c0e9f5fd43 100644 --- a/src/torchmetrics/functional/clustering/normalized_mutual_info_score.py +++ b/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 diff --git a/tests/unittests/clustering/test_adjusted_mutual_info_score.py b/tests/unittests/clustering/test_adjusted_mutual_info_score.py index 6e04e931cb5..079cf8d2ba5 100644 --- a/tests/unittests/clustering/test_adjusted_mutual_info_score.py +++ b/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.""" diff --git a/tests/unittests/clustering/test_normalized_mutual_info_score.py b/tests/unittests/clustering/test_normalized_mutual_info_score.py index 3e4140723f5..c965a6763ba 100644 --- a/tests/unittests/clustering/test_normalized_mutual_info_score.py +++ b/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."""