pyg-nightly 2.6.0.dev20240318__py3-none-any.whl → 2.7.0.dev20250115__py3-none-any.whl

torch_geometric/metrics/link_pred.py

@@ -1,4 +1,5 @@
-from typing import Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -14,6 +15,76 @@ except Exception:
     BaseMetric = torch.nn.Module  # type: ignore
 
 
+@dataclass(repr=False)
+class LinkPredMetricData:
+    pred_index_mat: Tensor
+    edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]]
+    edge_label_weight: Optional[Tensor] = None
+
+    @property
+    def pred_rel_mat(self) -> Tensor:
+        r"""Returns a matrix indicating the relevance of the `k`-th prediction.
+        If :obj:`edge_label_weight` is not given, relevance will be denoted as
+        binary.
+        """
+        if hasattr(self, '_pred_rel_mat'):
+            return self._pred_rel_mat  # type: ignore
+
+        # Flatten both prediction and ground-truth indices, and determine
+        # overlaps afterwards via `torch.searchsorted`.
+        max_index = max(  # type: ignore
+            self.pred_index_mat.max()
+            if self.pred_index_mat.numel() > 0 else 0,
+            self.edge_label_index[1].max()
+            if self.edge_label_index[1].numel() > 0 else 0,
+        ) + 1
+        arange = torch.arange(
+            start=0,
+            end=max_index * self.pred_index_mat.size(0),  # type: ignore
+            step=max_index,  # type: ignore
+            device=self.pred_index_mat.device,
+        ).view(-1, 1)
+        flat_pred_index = (self.pred_index_mat + arange).view(-1)
+        flat_label_index = max_index * self.edge_label_index[0]
+        flat_label_index = flat_label_index + self.edge_label_index[1]
+        flat_label_index, perm = flat_label_index.sort()
+        edge_label_weight = self.edge_label_weight
+        if edge_label_weight is not None:
+            assert edge_label_weight.size() == self.edge_label_index[0].size()
+            edge_label_weight = edge_label_weight[perm]
+
+        pos = torch.searchsorted(flat_label_index, flat_pred_index)
+        pos = pos.clamp(max=flat_label_index.size(0) - 1)  # Out-of-bounds.
+
+        pred_rel_mat = flat_label_index[pos] == flat_pred_index  # Find matches
+        if edge_label_weight is not None:
+            pred_rel_mat = edge_label_weight[pos].where(
+                pred_rel_mat,
+                pred_rel_mat.new_zeros(1),
+            )
+        pred_rel_mat = pred_rel_mat.view(self.pred_index_mat.size())
+
+        self._pred_rel_mat = pred_rel_mat
+        return pred_rel_mat
+
+    @property
+    def label_count(self) -> Tensor:
+        r"""The number of ground-truth labels for every example."""
+        if hasattr(self, '_label_count'):
+            return self._label_count  # type: ignore
+
+        label_count = scatter(
+            torch.ones_like(self.edge_label_index[0]),
+            self.edge_label_index[0],
+            dim=0,
+            dim_size=self.pred_index_mat.size(0),
+            reduce='sum',
+        )
+
+        self._label_count = label_count
+        return label_count
+
+
 class LinkPredMetric(BaseMetric):
     r"""An abstract class for computing link prediction retrieval metrics.
 
@@ -23,6 +94,7 @@ class LinkPredMetric(BaseMetric):
     is_differentiable: bool = False
     full_state_update: bool = False
     higher_is_better: Optional[bool] = None
+    weighted: bool = False
 
     def __init__(self, k: int) -> None:
         super().__init__()
@@ -47,6 +119,7 @@ class LinkPredMetric(BaseMetric):
         self,
         pred_index_mat: Tensor,
         edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
     ) -> None:
         r"""Updates the state variables based on the current mini-batch
         prediction.
@@ -62,45 +135,30 @@ class LinkPredMetric(BaseMetric):
             edge_label_index (torch.Tensor): The ground-truth indices for every
                 example in the mini-batch, given in COO format of shape
                 :obj:`[2, num_ground_truth_indices]`.
+            edge_label_weight (torch.Tensor, optional): The weight of the
+                ground-truth indices for every example in the mini-batch of
+                shape :obj:`[num_ground_truth_indices]`. If given, needs to be
+                a vector of positive values. Required for weighted metrics,
+                ignored otherwise. (default: :obj:`None`)
         """
-        if pred_index_mat.size(1) != self.k:
-            raise ValueError(f"Expected 'pred_index_mat' to hold {self.k} "
-                             f"many indices for every entry "
-                             f"(got {pred_index_mat.size(1)})")
-
-        # Compute a boolean matrix indicating if the k-th prediction is part of
-        # the ground-truth. We do this by flattening both prediction and
-        # target indices, and then determining overlaps via `torch.isin`.
-        max_index = max(  # type: ignore
-            pred_index_mat.max() if pred_index_mat.numel() > 0 else 0,
-            edge_label_index[1].max()
-            if edge_label_index[1].numel() > 0 else 0,
-        ) + 1
-        arange = torch.arange(
-            start=0,
-            end=max_index * pred_index_mat.size(0),
-            step=max_index,
-            device=pred_index_mat.device,
-        ).view(-1, 1)
-        flat_pred_index = (pred_index_mat + arange).view(-1)
-        flat_y_index = max_index * edge_label_index[0] + edge_label_index[1]
-
-        pred_isin_mat = torch.isin(flat_pred_index, flat_y_index)
-        pred_isin_mat = pred_isin_mat.view(pred_index_mat.size())
-
-        # Compute the number of targets per example:
-        y_count = scatter(
-            torch.ones_like(edge_label_index[0]),
-            edge_label_index[0],
-            dim=0,
-            dim_size=pred_index_mat.size(0),
-            reduce='sum',
+        if self.weighted and edge_label_weight is None:
+            raise ValueError(f"'edge_label_weight' is a required argument for "
+                             f"weighted '{self.__class__.__name__}' metrics")
+        if not self.weighted:
+            edge_label_weight = None
+
+        data = LinkPredMetricData(
+            pred_index_mat=pred_index_mat,
+            edge_label_index=edge_label_index,
+            edge_label_weight=edge_label_weight,
         )
+        self._update(data)
 
-        metric = self._compute(pred_isin_mat, y_count)
+    def _update(self, data: LinkPredMetricData) -> None:
+        metric = self._compute(data)
 
         self.accum += metric.sum()
-        self.total += (y_count > 0).sum()
+        self.total += (data.label_count > 0).sum()
 
     def compute(self) -> Tensor:
         r"""Computes the final metric value."""
@@ -109,28 +167,159 @@ class LinkPredMetric(BaseMetric):
         return self.accum / self.total
 
     def reset(self) -> None:
-        r"""Reset metric state variables to their default value."""
+        r"""Resets metric state variables to their default value."""
         if WITH_TORCHMETRICS:
             super().reset()
         else:
             self.accum.zero_()
             self.total.zero_()
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        r"""Compute the specific metric.
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        r"""Computes the specific metric.
         To be implemented separately for each metric class.
 
         Args:
-            pred_isin_mat (torch.Tensor): A boolean matrix whose :obj:`(i,k)`
-                element indicates if the :obj:`k`-th prediction for the
-                :obj:`i`-th example is correct or not.
-            y_count (torch.Tensor): A vector indicating the number of
-                ground-truth labels for each example.
+            data (LinkPredMetricData): The mini-batch data for computing a link
+                prediction metric per example.
         """
         raise NotImplementedError
 
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}(k={self.k})'
+        weighted_repr = ', weighted=True' if self.weighted else ''
+        return f'{self.__class__.__name__}(k={self.k}{weighted_repr})'
+
+
+class LinkPredMetricCollection(torch.nn.ModuleDict):
+    r"""A collection of metrics to reduce and speed-up computation of link
+    prediction metrics.
+
+    .. code-block:: python
+
+        from torch_geometric.metrics import (
+            LinkPredMAP,
+            LinkPredMetricCollection,
+            LinkPredPrecision,
+            LinkPredRecall,
+        )
+
+        metrics = LinkPredMetricCollection([
+            LinkPredMAP(k=10),
+            LinkPredPrecision(k=100),
+            LinkPredRecall(k=50),
+        ])
+
+        metrics.update(pred_index_mat, edge_label_index)
+        out = metrics.compute()
+        metrics.reset()
+
+        print(out)
+        >>> {'LinkPredMAP@10': tensor(0.375),
+        ...  'LinkPredPrecision@100': tensor(0.127),
+        ...  'LinkPredRecall@50': tensor(0.483)}
+
+    Args:
+        metrics: The link prediction metrics.
+    """
+    def __init__(
+        self,
+        metrics: Union[
+            List[LinkPredMetric],
+            Dict[str, LinkPredMetric],
+        ],
+    ) -> None:
+        super().__init__()
+
+        if isinstance(metrics, (list, tuple)):
+            metrics = {
+                f'{metric.__class__.__name__}@{metric.k}': metric
+                for metric in metrics
+            }
+        assert len(metrics) > 0
+        assert isinstance(metrics, dict)
+
+        for name, metric in metrics.items():
+            self[name] = metric
+
+    @property
+    def max_k(self) -> int:
+        r"""The maximum number of top-:math:`k` predictions to evaluate
+        against.
+        """
+        return max([metric.k for metric in self.values()])
+
+    @property
+    def weighted(self) -> bool:
+        r"""Returns :obj:`True` in case the collection holds at least one
+        weighted link prediction metric.
+        """
+        return any([metric.weighted for metric in self.values()])
+
+    def update(  # type: ignore
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        r"""Updates the state variables based on the current mini-batch
+        prediction.
+
+        :meth:`update` can be repeated multiple times to accumulate the results
+        of successive predictions, *e.g.*, inside a mini-batch training or
+        evaluation loop.
+
+        Args:
+            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
+                every example in the mini-batch with shape
+                :obj:`[batch_size, k]`.
+            edge_label_index (torch.Tensor): The ground-truth indices for every
+                example in the mini-batch, given in COO format of shape
+                :obj:`[2, num_ground_truth_indices]`.
+            edge_label_weight (torch.Tensor, optional): The weight of the
+                ground-truth indices for every example in the mini-batch of
+                shape :obj:`[num_ground_truth_indices]`. If given, needs to be
+                a vector of positive values. Required for weighted metrics,
+                ignored otherwise. (default: :obj:`None`)
+        """
+        if self.weighted and edge_label_weight is None:
+            raise ValueError(f"'edge_label_weight' is a required argument for "
+                             f"weighted '{self.__class__.__name__}' metrics")
+        if not self.weighted:
+            edge_label_weight = None
+
+        data = LinkPredMetricData(  # Share metric data across metrics.
+            pred_index_mat=pred_index_mat,
+            edge_label_index=edge_label_index,
+            edge_label_weight=edge_label_weight,
+        )
+
+        for metric in self.values():
+            if metric.weighted:
+                metric._update(data)
+                if WITH_TORCHMETRICS:
+                    metric._update_count += 1
+
+        data.edge_label_weight = None
+        if hasattr(data, '_pred_rel_mat'):
+            data._pred_rel_mat = data._pred_rel_mat != 0.0
+
+        for metric in self.values():
+            if not metric.weighted:
+                metric._update(data)
+                if WITH_TORCHMETRICS:
+                    metric._update_count += 1
+
+    def compute(self) -> Dict[str, Tensor]:
+        r"""Computes the final metric values."""
+        return {name: metric.compute() for name, metric in self.items()}
+
+    def reset(self) -> None:
+        r"""Reset metric state variables to their default value."""
+        for metric in self.values():
+            metric.reset()
+
+    def __repr__(self) -> str:
+        names = [f'  {name}: {metric},\n' for name, metric in self.items()]
+        return f'{self.__class__.__name__}([\n{"".join(names)}])'
 
 
 class LinkPredPrecision(LinkPredMetric):
@@ -140,9 +329,11 @@ class LinkPredPrecision(LinkPredMetric):
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        return pred_isin_mat.sum(dim=-1) / self.k
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        return pred_rel_mat.sum(dim=-1) / self.k
 
 
 class LinkPredRecall(LinkPredMetric):
@@ -152,9 +343,11 @@ class LinkPredRecall(LinkPredMetric):
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        return pred_isin_mat.sum(dim=-1) / y_count.clamp(min=1e-7)
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        return pred_rel_mat.sum(dim=-1) / data.label_count.clamp(min=1e-7)
 
 
 class LinkPredF1(LinkPredMetric):
@@ -164,11 +357,13 @@ class LinkPredF1(LinkPredMetric):
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        isin_count = pred_isin_mat.sum(dim=-1)
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        isin_count = pred_rel_mat.sum(dim=-1)
         precision = isin_count / self.k
-        recall = isin_count = isin_count / y_count.clamp(min=1e-7)
+        recall = isin_count / data.label_count.clamp(min=1e-7)
         return 2 * precision * recall / (precision + recall).clamp(min=1e-7)
 
 
@@ -180,12 +375,15 @@ class LinkPredMAP(LinkPredMetric):
         k (int): The number of top-:math:`k` predictions to evaluate against.
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        cum_precision = (torch.cumsum(pred_isin_mat, dim=1) /
-                         torch.arange(1, self.k + 1, device=y_count.device))
-        return ((cum_precision * pred_isin_mat).sum(dim=-1) /
-                y_count.clamp(min=1e-7, max=self.k))
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        device = pred_rel_mat.device
+        arange = torch.arange(1, pred_rel_mat.size(1) + 1, device=device)
+        cum_precision = pred_rel_mat.cumsum(dim=1) / arange
+        return ((cum_precision * pred_rel_mat).sum(dim=-1) /
+                data.label_count.clamp(min=1e-7, max=self.k))
 
 
 class LinkPredNDCG(LinkPredMetric):
@@ -194,25 +392,79 @@ class LinkPredNDCG(LinkPredMetric):
 
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
+        weighted (bool, optional): If set to :obj:`True`, assumes sorted lists
+            of ground-truth items according to a relevance score as given by
+            :obj:`edge_label_weight`. (default: :obj:`False`)
     """
     higher_is_better: bool = True
+    weighted: bool = False
 
-    def __init__(self, k: int):
+    def __init__(self, k: int, weighted: bool = False):
         super().__init__(k=k)
+        self.weighted = weighted
 
         dtype = torch.get_default_dtype()
-        multiplier = 1.0 / torch.arange(2, k + 2, dtype=dtype).log2()
+        discount = torch.arange(2, k + 2, dtype=dtype).log2()
+
+        self.discount: Tensor
+        self.register_buffer('discount', discount)
 
-        self.multiplier: Tensor
-        self.register_buffer('multiplier', multiplier)
+        if not weighted:
+            self.register_buffer('idcg', cumsum(1.0 / discount))
+        else:
+            self.idcg = None
 
-        self.idcg: Tensor
-        self.register_buffer('idcg', cumsum(multiplier))
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        discount = self.discount[:pred_rel_mat.size(1)].view(1, -1)
+        dcg = (pred_rel_mat / discount).sum(dim=-1)
 
-    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        dcg = (pred_isin_mat * self.multiplier.view(1, -1)).sum(dim=-1)
-        idcg = self.idcg[y_count.clamp(max=self.k)]
+        if not self.weighted:
+            assert self.idcg is not None
+            idcg = self.idcg[data.label_count.clamp(max=self.k)]
+        else:
+            assert data.edge_label_weight is not None
+            # Sort weights within example-wise buckets via two sorts to get the
+            # local index order within buckets:
+            weight, batch = data.edge_label_weight, data.edge_label_index[0]
+            perm1 = weight.argsort(descending=True)
+            perm2 = batch[perm1].argsort(stable=True)
+            global_index = torch.empty_like(perm1)
+            global_index[perm1[perm2]] = torch.arange(
+                global_index.size(0), device=global_index.device)
+            local_index = global_index - cumsum(data.label_count)[batch]
+
+            # Get the discount per local index:
+            discount = torch.cat([
+                self.discount,
+                self.discount.new_full((1, ), fill_value=float('inf')),
+            ])
+            discount = discount[local_index.clamp(max=self.k + 1)]
+
+            idcg = scatter(  # Apply discount and aggregate:
+                weight / discount,
+                batch,
+                dim_size=data.pred_index_mat.size(0),
+                reduce='sum',
+            )
 
         out = dcg / idcg
         out[out.isnan() | out.isinf()] = 0.0
         return out
+
+
+class LinkPredMRR(LinkPredMetric):
+    r"""A link prediction metric to compute the MRR @ :math:`k` (Mean
+    Reciprocal Rank).
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    higher_is_better: bool = True
+    weighted: bool = False
+
+    def _compute(self, data: LinkPredMetricData) -> Tensor:
+        pred_rel_mat = data.pred_rel_mat[:, :self.k]
+        device = pred_rel_mat.device
+        arange = torch.arange(1, pred_rel_mat.size(1) + 1, device=device)
+        return (pred_rel_mat / arange).max(dim=-1)[0]

torch_geometric/nn/aggr/__init__.py

@@ -24,6 +24,8 @@ from .mlp import MLPAggregation
 from .deep_sets import DeepSetsAggregation
 from .set_transformer import SetTransformerAggregation
 from .lcm import LCMAggregation
+from .variance_preserving import VariancePreservingAggregation
+from .patch_transformer import PatchTransformerAggregation
 
 __all__ = classes = [
     'Aggregation',
@@ -51,4 +53,6 @@ __all__ = classes = [
     'DeepSetsAggregation',
     'SetTransformerAggregation',
     'LCMAggregation',
+    'VariancePreservingAggregation',
+    'PatchTransformerAggregation',
 ]

torch_geometric/nn/aggr/attention.py

@@ -65,8 +65,6 @@ class AttentionalAggregation(Aggregation):
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
 
-        self.assert_two_dimensional_input(x, dim)
-
         if self.gate_mlp is not None:
             gate = self.gate_mlp(x, batch=index, batch_size=dim_size)
         else:

torch_geometric/nn/aggr/base.py

@@ -25,7 +25,7 @@ class Aggregation(torch.nn.Module):
     Notably, :obj:`index` does not have to be sorted (for most aggregation
     operators):
 
-    .. code-block::
+    .. code-block:: python
 
        # Feature matrix holding 10 elements with 64 features each:
        x = torch.randn(10, 64)
@@ -39,7 +39,7 @@ class Aggregation(torch.nn.Module):
     called :obj:`ptr`. Here, elements within the same set need to be grouped
     together in the input, and :obj:`ptr` defines their boundaries:
 
-    .. code-block::
+    .. code-block:: python
 
        # Feature matrix holding 10 elements with 64 features each:
        x = torch.randn(10, 64)
@@ -47,7 +47,7 @@ class Aggregation(torch.nn.Module):
        # Define the boundary indices for three sets:
        ptr = torch.tensor([0, 4, 7, 10])
 
-       output = aggr(x, ptr=ptr)  #  Output shape: [4, 64]
+       output = aggr(x, ptr=ptr)  #  Output shape: [3, 64]
 
     Note that at least one of :obj:`index` or :obj:`ptr` must be defined.
 
@@ -94,11 +94,9 @@ class Aggregation(torch.nn.Module):
             max_num_elements: (int, optional): The maximum number of elements
                 within a single aggregation group. (default: :obj:`None`)
         """
-        pass
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        pass
 
     @disable_dynamic_shapes(required_args=['dim_size'])
     def __call__(