Mutual Information Score

docs/source/clustering/mutual_info_score.rst

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

docs/source/index.rst

@@ -158,6 +158,14 @@ Or directly from conda
 
    classification/*
 
+.. toctree::
+   :maxdepth: 2
+   :name: clustering
+   :caption: Clustering
+   :glob:
+
+   clustering/*
+
 .. toctree::
    :maxdepth: 2
    :name: detection

src/torchmetrics/clustering/__init__.py

@@ -0,0 +1,18 @@
+# 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 torchmetrics.clustering.mutual_info_score import MutualInfoScore
+
+__all__ = [
+    "MutualInfoScore",
+]

src/torchmetrics/clustering/mutual_info_score.py

@@ -0,0 +1,124 @@
+# 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, Optional, Sequence, Union
+
+import torch
+from torch import Tensor
+
+from torchmetrics.functional.clustering.mutual_info_score import _mutual_info_score_compute, _mutual_info_score_update
+from torchmetrics.metric import Metric
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
+
+if not _MATPLOTLIB_AVAILABLE:
+    __doctest_skip__ = ["MutualInfoScore.plot"]
+
+
+class MutualInfoScore(Metric):
+    r"""Compute `Mutual Information Score`_.
+
+    .. math::
+        MI(U,V) = \sum_{i=1}^{\abs{U}} \sum_{j=1}^{\abs{V}} \frac{\abs{U_i\cap V_j}}{N} \log\frac{N\abs{U_i\cap V_j}}{\abs{U_i}\abs{V_j}}
+
+    Where :math:`U` is a tensor of target values, :math:`V` is a tensor of predictions,
+    :math:`\abs{U_i}` is the number of samples in cluster :math:`U_i`, and
+    :math:`\abs{V_i}` is the number of samples in cluster :math:`V_i`.
+
+    The metric is symmetric, therefore swapping :math:`U` and :math:`V` yields
+    the same mutual information score.
+
+    Args:
+        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+
+    As input to ``forward`` and ``update`` the metric accepts the following input:
+
+    - ``preds`` (:class:`~torch.Tensor`): either single output float tensor with shape ``(N,)``
+    - ``target`` (:class:`~torch.Tensor`): either single output tensor with shape ``(N,)``
+
+    As output of ``forward`` and ``compute`` the metric returns the following output:
+
+    - ``mi_score`` (:class:`~torch.Tensor`): A tensor with the Mutual Information Score
+
+    Example:
+        >>> from torchmetrics.clustering import MutualInfoScore
+        >>> target = torch.tensor([])
+        >>> preds = torch.tensor([])
+        >>> mi_score = MutualInfoScore()
+        >>> mi_score(preds, target)
+        tensor()
+
+    """
+
+    is_differentiable = True
+    higher_is_better = None
+    full_state_update: bool = True
+    plot_lower_bound: float = 0.0
+    plot_upper_bound: float = 1.0  # theoretical upper bound is +inf
+    preds: List[Tensor]
+    target: List[Tensor]
+    contingency: Tensor
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        # self.num_classes = num_classes
+        #
+        # self.add_state("contingency", default=torch.zeros(self.num_classes), dist_reduce_fx=None)
+
+    def update(self, preds: Tensor, target: Tensor) -> None:
+        """Update state with predictions and targets."""
+        self.contingency = _mutual_info_score_update(preds, target)
+
+    def compute(self) -> Tensor:
+        """Compute mutual information over state."""
+        return _mutual_info_score_compute(self.contingency)
+
+    def plot(self, val: Union[Tensor, Sequence[Tensor], None] = None, ax: Optional[_AX_TYPE] = None) -> _PLOT_OUT_TYPE:
+        """Plot a single or multiple values from the metric.
+
+        Args:
+            val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
+                If no value is provided, will automatically call `metric.compute` and plot that result.
+            ax: An matplotlib axis object. If provided will add plot to that axis
+
+        Returns:
+            Figure and Axes object
+
+        Raises:
+            ModuleNotFoundError:
+                If `matplotlib` is not installed
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting a single value
+            >>> import torch
+            >>> from torchmetrics.clustering import MutualInfoScore
+            >>> metric = MutualInfoScore(num_classes=5)
+            >>> metric.update(torch.randint(0, 4, (100,)), torch.randint(0, 4, (100,)))
+            >>> fig_, ax_ = metric.plot()
+
+        .. plot::
+            :scale: 75
+
+            >>> # Example plotting multiple values
+            >>> import torch
+            >>> from torchmetrics.clustering import MutualInfoScore
+            >>> metric = MutualInfoScore(num_classes=5)
+            >>> values = [ ]
+            >>> for _ in range(10):
+            ...     values.append(metric(torch.randint(0, 4, (100,)), torch.randint(0, 4, (100,))))
+            >>> fig_, ax_ = metric.plot(values)
+
+        """
+        return self._plot(val, ax)

src/torchmetrics/functional/clustering/__init__.py

@@ -0,0 +1,3 @@
+from torchmetrics.functional.clustering.mutual_info_score import mutual_info_score
+
+__all__ = ["mutual_info_score"]

src/torchmetrics/functional/clustering/mutual_info_score.py

@@ -0,0 +1,130 @@
+# 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 Optional, Tuple
+
+import torch
+from torch import Tensor, tensor
+
+from torchmetrics.utilities.checks import _check_same_shape
+
+
+def _check_cluster_labels(preds: Tensor, target: Tensor) -> None:
+    """Check shape of input tensors."""
+    _check_same_shape(preds, target)
+    if (
+        torch.is_floating_point(preds)
+        or torch.is_complex(preds)
+        or torch.is_floating_point(target)
+        or torch.is_complex(target)
+    ):
+        raise ValueError(
+            f"Expected real, discrete values but received {preds.dtype} for"
+            f"predictions and {target.dtype} for target labels instead."
+        )
+
+
+def _calculate_contingency_matrix(
+    preds: Tensor, target: Tensor, eps: Optional[float] = 1e-16, sparse: bool = False
+) -> Tensor:
+    """Calculate contingency matrix.
+
+    Args:
+        preds: predicted labels
+        target: ground truth labels
+        sparse: If True, returns contingency matrix as a sparse matrix.
+
+    Returns:
+        contingency: contingency matrix of shape (n_classes_target, n_classes_preds)
+
+    """
+    if eps is not None and sparse is True:
+        raise ValueError("Cannot specify `eps` and return sparse tensor.")
+
+    preds_classes, preds_idx = torch.unique(preds, return_inverse=True)
+    target_classes, target_idx = torch.unique(target, return_inverse=True)
+
+    n_classes_preds = preds_classes.size(0)
+    n_classes_target = target_classes.size(0)
+
+    contingency = torch.sparse_coo_tensor(
+        torch.stack((target_idx, preds_idx)), torch.ones(target_idx.size(0)), (n_classes_target, n_classes_preds)
+    )
+
+    if not sparse:
+        contingency = contingency.to_dense()
+    if eps:
+        contingency = contingency + eps
+
+    return contingency
+
+
+def _mutual_info_score_update(
+    preds: Tensor,
+    target: Tensor,
+    # num_classes: int
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """Update and return variables required to compute the mutual information score.
+
+    Args:
+        preds: predicted class labels
+        target: ground truth class labels
+
+    Returns:
+        contingency: contingency matrix
+
+    """
+    _check_cluster_labels(preds, target)
+    return _calculate_contingency_matrix(preds, target)
+
+
+def _mutual_info_score_compute(contingency: Tensor) -> Tensor:
+    """Compute the mutual information score based on the contingency matrix.
+
+    Args:
+        contingency: contingency matrix
+
+    Returns:
+        mutual_info: mutual information score
+
+    """
+    N = contingency.sum()
+    U = contingency.sum(dim=1)
+    V = contingency.sum(dim=0)
+
+    # Check if preds or target labels only have one cluster
+    if U.size() == 1 or V.size() == 1:
+        return tensor(0.0)
+
+    log_outer = torch.log(U).reshape(-1, 1) + torch.log(V)
+    mutual_info = contingency / N * (torch.log(N) + torch.log(contingency) - log_outer)
+    return mutual_info.sum()
+
+
+def mutual_info_score(preds: Tensor, target: Tensor) -> Tensor:
+    """Compute mutual information between two clusterings.
+
+    Args:
+        preds: predicted classes
+        target: ground truth classes
+
+    Example:
+        >>> from torchmetrics.functional.clustering import mutual_info_score
+        >>> target = torch.tensor([0, 3, 2, 2, 1])
+        >>> preds = torch.tensor([1, 3, 2, 0, 1])
+        >>> mutual_info_score(preds, target)
+        tensor([1.05492])
+
+    """
+    contingency = _mutual_info_score_update(preds, target)
+    return _mutual_info_score_compute(contingency)