from typing import Optional
from torch import Tensor
from torch_geometric.typing import Adj
from tgp.connect import SparseConnect
from tgp.lift import BaseLift
from tgp.reduce import BaseReduce
from tgp.select import GraclusSelect, SelectOutput
from tgp.src import BasePrecoarseningMixin, PoolingOutput, SRCPooling
from tgp.utils.typing import ConnectionType, LiftType, ReduceType, SinvType
[docs]
class GraclusPooling(BasePrecoarseningMixin, SRCPooling):
r"""The Graclus pooling operator inspired by the paper `"Weighted Graph Cuts without
Eigenvectors: A Multilevel Approach" <https://ieeexplore.ieee.org/document/4302760>`_
(Dhillon et al., TPAMI 2007).
+ The :math:`\texttt{select}` operator is implemented with :class:`~tgp.select.GraclusSelect`.
+ The :math:`\texttt{reduce}` operator is implemented with :class:`~tgp.reduce.BaseReduce`.
+ The :math:`\texttt{connect}` operator is implemented with :class:`~tgp.connect.SparseConnect`.
+ The :math:`\texttt{lift}` operator is implemented with :class:`~tgp.lift.BaseLift`.
Args:
lift (~tgp.utils.typing.LiftType, optional):
Defines how to compute the matrix :math:`\mathbf{S}_\text{inv}` to lift the pooled node features.
- ``"precomputed"`` (default): Use as :math:`\mathbf{S}_\text{inv}` what is
already stored in the ``"s_inv"`` attribute of the :class:`~tgp.select.SelectOutput`.
- ``"transpose"``: Recomputes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^\top`,
the transpose of :math:`\mathbf{S}`.
- ``"inverse"``: Recomputes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^+`,
the Moore-Penrose pseudoinverse of :math:`\mathbf{S}`.
s_inv_op (~tgp.utils.typing.SinvType, optional):
The operation used to compute :math:`\mathbf{S}_\text{inv}` from the select matrix
:math:`\mathbf{S}`. :math:`\mathbf{S}_\text{inv}` is stored in the ``"s_inv"`` attribute of
the :class:`~tgp.select.SelectOutput`. It can be one of:
- ``"transpose"`` (default): Computes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^\top`,
the transpose of :math:`\mathbf{S}`.
- ``"inverse"``: Computes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^+`,
the Moore-Penrose pseudoinverse of :math:`\mathbf{S}`.
lift_red_op (~tgp.utils.typing.ReduceType, optional):
The aggregation function to be applied to the lifted node features.
Can be any string of class :class:`~tgp.utils.typing.ReduceType` admitted by
:obj:`~torch_geometric.utils.scatter`,
e.g., ``'sum'``, ``'mean'``, ``'max'``)
(default: ``"sum"``)
cached (bool, optional):
If set to :obj:`True`, the output of the :math:`\texttt{select}` and :math:`\texttt{select}`
operations will be cached, so that they do not need to be recomputed.
(default: :obj:`False`)
remove_self_loops (bool, optional):
If :obj:`True`, the self-loops will be removed from the adjacency matrix.
(default: :obj:`True`)
degree_norm (bool, optional):
If :obj:`True`, the adjacency matrix will be symmetrically normalized.
(default: :obj:`False`)
edge_weight_norm (bool, optional):
Whether to normalize the edge weights by dividing by the maximum absolute value per graph.
(default: :obj:`False`)
"""
def __init__(
self,
lift: LiftType = "precomputed",
s_inv_op: SinvType = "transpose",
connect_red_op: ConnectionType = "sum",
lift_red_op: ReduceType = "sum",
cached: bool = False,
remove_self_loops: bool = True,
degree_norm: bool = False,
edge_weight_norm: bool = False,
):
super().__init__(
selector=GraclusSelect(s_inv_op=s_inv_op),
reducer=BaseReduce(),
lifter=BaseLift(matrix_op=lift, reduce_op=lift_red_op),
connector=SparseConnect(
reduce_op=connect_red_op,
remove_self_loops=remove_self_loops,
degree_norm=degree_norm,
edge_weight_norm=edge_weight_norm,
),
cached=cached,
)
self.cached = cached
[docs]
def forward(
self,
x: Tensor,
adj: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
so: Optional[SelectOutput] = None,
batch: Optional[Tensor] = None,
lifting: bool = False,
**kwargs,
) -> PoolingOutput:
r"""Forward pass.
Args:
x (~torch.Tensor): The node feature matrix of shape :math:`[N, F]`,
where :math:`N` is the number of nodes in the batch and
:math:`F` is the number of node features.
adj (~torch_geometric.typing.Adj, optional): The connectivity matrix.
It can either be a ``torch_sparse.SparseTensor`` of (sparse) shape :math:`[N, N]`,
where :math:`N` is the number of nodes in the batch or a :obj:`~torch.Tensor` of shape
:math:`[2, E]`, where :math:`E` is the number of edges in the batch.
If ``lifting`` is :obj:`False`, it cannot be :obj:`None`.
(default: :obj:`None`)
edge_weight (~torch.Tensor, optional): A vector of shape :math:`[E]` or :math:`[E, 1]`
containing the weights of the edges.
(default: :obj:`None`)
so (~tgp.select.SelectOutput, optional): The output of the :math:`\texttt{select}` operator.
(default: :obj:`None`)
batch (torch.Tensor, optional): The batch vector
:math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which indicates
to which graph in the batch each node belongs. (default: :obj:`None`)
lifting (bool, optional): If set to :obj:`True`, the :math:`\texttt{lift}` operation is performed.
(default: :obj:`False`)
Returns:
PoolingOutput: The output of the pooling operator.
"""
if lifting:
# Lift
x_lifted = self.lift(x_pool=x, so=so)
return x_lifted
else:
# Select
so = self.select(
edge_index=adj, edge_weight=edge_weight, num_nodes=x.size(0)
)
# Reduce
x_pooled, batch_pooled = self.reduce(x=x, so=so, batch=batch)
# Connect
edge_index_pooled, edge_weight_pooled = self.connect(
edge_index=adj,
so=so,
edge_weight=edge_weight,
batch_pooled=batch_pooled,
)
out = PoolingOutput(
x=x_pooled,
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
)
return out
def extra_repr_args(self) -> dict:
return {"cached": self.cached}
