New metric: Normalized Mutual Information Score

CHANGELOG.md

@@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `RandScore` metric to cluster package ([#2025](https://github.com/Lightning-AI/torchmetrics/pull/2025)
 
 
+- Added `NormalizedMutualInfoScore` metric to cluster package ([#2029](https://github.com/Lightning-AI/torchmetrics/pull/2029)
+
+
 ### Changed
 
 -

docs/source/clustering/normalized_mutual_info_score.rst

@@ -0,0 +1,21 @@
+.. customcarditem::
+   :header: Normalized Mutual Information Score
+   :image: https://pl-flash-data.s3.amazonaws.com/assets/thumbnails/clustering.svg
+   :tags: Clustering
+
+.. include:: ../links.rst
+
+###################################
+Normalized Mutual Information Score
+###################################
+
+Module Interface
+________________
+
+.. autoclass:: torchmetrics.clustering.NormalizedMutualInfoScore
+    :exclude-members: update, compute
+
+Functional Interface
+____________________
+
+.. autofunction:: torchmetrics.functional.clustering.normalized_mutual_info_score

docs/source/links.rst

@@ -151,5 +151,6 @@
 .. _DIOU: https://arxiv.org/abs/1911.08287v1
 .. _GIOU: https://arxiv.org/abs/1902.09630
 .. _Mutual Information Score: https://en.wikipedia.org/wiki/Mutual_information
+.. _Normalized Mutual Information Score: https://scikit-learn.org/stable/modules/generated/sklearn.metrics.normalized_mutual_info_score.html
 .. _pycocotools: https://github.com/cocodataset/cocoapi/tree/master/PythonAPI/pycocotools
 .. _Rand Score: https://link.springer.com/article/10.1007/BF01908075

src/torchmetrics/clustering/__init__.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from torchmetrics.clustering.mutual_info_score import MutualInfoScore
+from torchmetrics.clustering.normalized_mutual_info_score import NormalizedMutualInfoScore
 from torchmetrics.clustering.rand_score import RandScore
 
 __all__ = [
     "MutualInfoScore",
+    "NormalizedMutualInfoScore",
     "RandScore",
 ]

src/torchmetrics/clustering/normalized_mutual_info_score.py

@@ -0,0 +1,80 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, List, Literal
+
+from torch import Tensor
+
+from torchmetrics.clustering.mutual_info_score import MutualInfoScore
+from torchmetrics.functional.clustering.normalized_mutual_info_score import normalized_mutual_info_score
+from torchmetrics.utilities.data import dim_zero_cat
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
+
+if not _MATPLOTLIB_AVAILABLE:
+    __doctest_skip__ = ["NormalizedMutualInfoScore.plot"]
+
+
+class NormalizedMutualInfoScore(MutualInfoScore):
+    r"""Compute `Normalized Mutual Information Score`_.
+
+    .. math::
+        NMI(U,V) = \frac{MI(U,V)}{M_p(U,V)}
+
+    Where :math:`U` is a tensor of target values, :math:`V` is a tensor of predictions,
+    :math:`M_p(U,V)` is the generalized mean of order :math:`p` of :math:`U` and :math:`V`,
+    and :math:`MI(U,V)` is the mutual information score between clusters :math:`U` and :math:`V`.
+
+    The metric is symmetric, therefore swapping :math:`U` and :math:`V` yields
+    the same mutual information score.
+
+    As input to ``forward`` and ``update`` the metric accepts the following input:
+
+    - ``preds`` (:class:`~torch.Tensor`): single integer tensor with shape ``(N,)`` with predicted cluster labels
+    - ``target`` (:class:`~torch.Tensor`): single integer tensor with shape ``(N,)`` with ground truth cluster labels
+
+    As output of ``forward`` and ``compute`` the metric returns the following output:
+
+    - ``nmi_score`` (:class:`~torch.Tensor`): A tensor with the Normalized Mutual Information Score
+
+    Args:
+        average_method: Method used to calculate generalized mean for normalization
+        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+
+    Example:
+        >>> import torch
+        >>> from torchmetrics.clustering import NormalizedMutualInfoScore
+        >>> preds = torch.tensor([2, 1, 0, 1, 0])
+        >>> target = torch.tensor([0, 2, 1, 1, 0])
+        >>> nmi_score = NormalizedMutualInfoScore("arithmetic")
+        >>> nmi_score(preds, target)
+        tensor(0.4744)
+
+    """
+
+    is_differentiable = True
+    higher_is_better = None
+    full_state_update: bool = True
+    plot_lower_bound: float = 0.0
+    plot_upper_bound: float = 0.0
+    preds: List[Tensor]
+    target: List[Tensor]
+    contingency: Tensor
+
+    def __init__(self, average_method: Literal["min", "geometric", "arithmetic", "max"] = "arithmetic", **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+        self.average_method = average_method
+
+    def compute(self) -> Tensor:
+        """Compute normalized mutual information over state."""
+        return normalized_mutual_info_score(dim_zero_cat(self.preds), dim_zero_cat(self.target), self.average_method)

src/torchmetrics/functional/clustering/__init__.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from torchmetrics.functional.clustering.mutual_info_score import mutual_info_score
+from torchmetrics.functional.clustering.normalized_mutual_info_score import normalized_mutual_info_score
 from torchmetrics.functional.clustering.rand_score import rand_score
 
-__all__ = ["mutual_info_score", "rand_score"]
+__all__ = [
+    "mutual_info_score",
+    "normalized_mutual_info_score",
+    "rand_score",
+]

src/torchmetrics/functional/clustering/normalized_mutual_info_score.py

@@ -0,0 +1,52 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from torchmetrics.functional.clustering.mutual_info_score import mutual_info_score
+from torchmetrics.functional.clustering.utils import calculate_entropy, calculate_generalized_mean
+
+
+def normalized_mutual_info_score(
+    preds: Tensor, target: Tensor, average_method: Literal["min", "geometric", "arithmetic", "max"] = "arithmetic"
+) -> Tensor:
+    """Compute normalized mutual information between two clusterings.
+
+    Args:
+        preds: predicted cluster labels
+        target: ground truth cluster labels
+        average_method: normalizer computation method
+
+    Returns:
+        Scalar tensor with normalized mutual info score between 0.0 and 1.0
+
+    Example:
+        >>> from torchmetrics.functional.clustering import normalized_mutual_info_score
+        >>> target = torch.tensor([0, 3, 2, 2, 1])
+        >>> preds = torch.tensor([1, 3, 2, 0, 1])
+        >>> normalized_mutual_info_score(preds, target, "arithmetic")
+        tensor(0.7919)
+
+    """
+    mutual_info = mutual_info_score(preds, target)
+    if torch.allclose(mutual_info, torch.tensor(0.0), atol=torch.finfo().eps):
+        return torch.tensor(0.0)
+
+    normalizer = calculate_generalized_mean(
+        torch.stack([calculate_entropy(preds), calculate_entropy(target)]), average_method
+    )
+
+    return mutual_info / normalizer

src/torchmetrics/functional/clustering/utils.py

@@ -11,14 +11,86 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional
+from typing import Optional, Union
 
 import torch
-from torch import Tensor
+from torch import Tensor, tensor
 
 from torchmetrics.utilities.checks import _check_same_shape
 
 
+def calculate_entropy(x: Tensor) -> Tensor:
+    """Calculate entropy for a tensor of labels.
+
+    Final calculation of entropy is performed in log form to account for roundoff error.
+
+    Args:
+        x: labels
+
+    Returns:
+        entropy: entropy of tensor
+
+    Example:
+        >>> from torchmetrics.functional.clustering.utils import calculate_entropy
+        >>> labels = torch.tensor([1, 3, 2, 2, 1])
+        >>> calculate_entropy(labels)
+        tensor(1.0549)
+
+    """
+    if len(x) == 0:
+        return tensor(1.0, device=x.device)
+
+    p = torch.bincount(torch.unique(x, return_inverse=True)[1])
+    p = p[p > 0]
+
+    if p.size() == 1:
+        return tensor(0.0, device=x.device)
+
+    n = p.sum()
+    return -torch.sum((p / n) * (torch.log(p) - torch.log(n)))
+
+
+def calculate_generalized_mean(x: Tensor, p: Union[int, Literal["min", "geometric", "arithmetic", "max"]) -> Tensor:
+    """Return generalized (power) mean of a tensor.
+
+    Args:
+        x: tensor
+        p: power
+
+    Returns:
+        generalized_mean: generalized mean
+
+    Example (p="min"):
+        >>> from torchmetrics.functional.clustering.utils import calculate_generalized_mean
+        >>> x = torch.tensor([1, 3, 2, 2, 1])
+        >>> calculate_generalized_mean(x, "min")
+        tensor(1)
+
+    Example (p="geometric"):
+        >>> from torchmetrics.functional.clustering.utils import calculate_generalized_mean
+        >>> x = torch.tensor([1, 3, 2, 2, 1])
+        >>> calculate_generalized_mean(x, "geometric")
+        tensor(1.6438)
+
+    """
+    if torch.is_complex(x) or torch.any(x <= 0.0):
+        raise ValueError("`x` must contain positive real numbers")
+
+    if isinstance(p, str):
+        if p == "min":
+            return x.min()
+        if p == "geometric":
+            return torch.exp(torch.mean(x.log()))
+        if p == "arithmetic":
+            return x.mean()
+        if p == "max":
+            return x.max()
+
+        raise ValueError("'method' must be 'min', 'geometric', 'arirthmetic', or 'max'")
+
+    return torch.mean(torch.pow(x, p)) ** (1.0 / p)
+
+
 def calculate_contingency_matrix(
     preds: Tensor, target: Tensor, eps: Optional[float] = None, sparse: bool = False
 ) -> Tensor: