Models¶

class CommonNeighbors(nn.Module):
    """
    Computes Common Neighbors scores for hyperedges.

    Attributes:
        scorer: An instance of a NeighborScorer that computes the scores for hyperedges.
    """

    def __init__(
        self,
        aggregation: Literal["mean", "min", "sum"],
        scorer: NeighborScorer | None = None,
    ) -> None:
        """
        Initialize the common-neighbors model.

        Args:
            aggregation: Method used by the default scorer to aggregate pairwise counts.
                Defaults to ``mean``.
            scorer: Optional custom neighbor scorer.
        """
        super().__init__()
        self.scorer: NeighborScorer = (
            scorer if scorer is not None else CommonNeighborsHyperedgeScorer(aggregation)
        )

    def forward(
        self,
        candidate_nodes: Tensor,
        hyperedge_index: Tensor | None = None,
        node_to_neighbors: dict[int, Neighborhood] | None = None,
    ) -> Tensor:
        """
        Compute CN scores for all hyperedges in the batch.

        Args:
            candidate_nodes: Tensor containing node IDs to score of shape ``(num_nodes,)``.
            hyperedge_index: Tensor containing the hyperedge indices. Defaults to ``None``.
            node_to_neighbors: Optional mapping from nodes to their neighborhoods.
                Defaults to ``None``.

        Returns:
            scores: A 1-D tensor of shape (num_hyperedges,) with CN scores.
        """
        scores = self.scorer.score_batch(
            candidate_nodes=candidate_nodes,
            hyperedge_index=hyperedge_index,
            node_to_neighbors=node_to_neighbors,
        )
        torch.log1p(scores, out=scores)
        return scores

class GCN(nn.Module):
    """
    A reusable multi-layer GCN stack built from ``torch_geometric.nn.GCNConv``.

    Attributes:
        dropout: Dropout layer applied after hidden GCN layers.
        activation: Activation module applied after hidden GCN layers.
        layers: GCN convolution layers.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int | None = None,
        num_layers: int = 2,
        drop_rate: float = 0.0,
        bias: bool = True,
        activation_fn: ActivationFn | None = None,
        activation_fn_kwargs: dict | None = None,
        improved: bool = False,
        add_self_loops: bool = True,
        normalize: bool = True,
        cached: bool = False,
    ):
        """
        Initialize the GCN stack.

        Args:
            in_channels: Dimension of input node embeddings.
            out_channels: Dimension of output node embeddings.
            hidden_channels: Dimension of hidden node embeddings.
            num_layers: Number of GCN layers. Must be at least ``1``.
                Defaults to ``2``.
            drop_rate: Dropout rate applied after hidden GCN layers.
                Defaults to ``0.0``.
            bias: Whether to include a bias term in GCN layers.
                Defaults to ``True``.
            activation_fn: Activation function class used after hidden layers.
                Defaults to ``nn.ReLU``.
            activation_fn_kwargs: Keyword arguments for the activation function.
                Defaults to empty dict.
            improved: Whether to use the improved version of ``GCNConv``.
            add_self_loops: Whether to add self-loops to the input graph.
            normalize: Whether to symmetrically normalize adjacency in ``GCNConv``.
            cached: Whether to cache the normalized adjacency matrix in ``GCNConv``.
                Defaults to ``False``.
        """
        super().__init__()
        activation_fn = activation_fn if activation_fn is not None else nn.ReLU
        activation_fn_kwargs = activation_fn_kwargs if activation_fn_kwargs is not None else {}

        self.dropout: nn.Dropout = nn.Dropout(drop_rate)
        self.activation: nn.Module = activation_fn(**activation_fn_kwargs)
        self.layers: nn.ModuleList = self.__build_layers(
            in_channels=in_channels,
            out_channels=out_channels,
            hidden_channels=hidden_channels,
            num_layers=num_layers,
            bias=bias,
            improved=improved,
            add_self_loops=add_self_loops,
            normalize=normalize,
            cached=cached,
        )

    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
        """
        Apply the GCN stack to node features.

        Args:
            x: Node feature matrix.
            edge_index: Graph edge index.

        Returns:
            x: Output node embeddings.
        """
        num_layers = len(self.layers)
        for idx, layer in enumerate(self.layers):
            x = layer(x, edge_index)

            is_not_last_layer = not is_layer(idx, num_layers - 1)
            if is_not_last_layer:
                x = self.activation(x)
                x = self.dropout(x)

        return x

    def __build_layers(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int | None,
        num_layers: int,
        bias: bool,
        improved: bool,
        add_self_loops: bool,
        normalize: bool,
        cached: bool,
    ) -> nn.ModuleList:
        """
        Build the internal GCN convolution layers.

        Args:
            in_channels: Dimension of input node embeddings.
            out_channels: Dimension of output node embeddings.
            hidden_channels: Dimension of hidden node embeddings.
            num_layers: Number of GCN layers.
            bias: Whether to include a bias term in GCN layers.
            improved: Whether to use the improved version of ``GCNConv``.
            add_self_loops: Whether to add self-loops to the input graph.
            normalize: Whether to symmetrically normalize adjacency in ``GCNConv``.
            cached: Whether ``GCNConv`` should cache normalized adjacency.

        Returns:
            layers: Module list containing the GCN layers.

        Raises:
            ValueError: If the layer configuration is invalid.
        """
        if num_layers < 1:
            raise ValueError(f"Expected num_layers >= 1 for GCN, got {num_layers}.")

        hidden_channels = hidden_channels if hidden_channels is not None else 0
        if num_layers > 1 and hidden_channels <= 0:
            raise ValueError(
                f"Expected positive hidden_channels for GCN with multiple layers, "
                f"got {hidden_channels}."
            )

        common_kwargs: dict[str, bool] = {
            "bias": bias,
            "improved": improved,
            "add_self_loops": add_self_loops,
            "normalize": normalize,
            "cached": cached,
        }

        if num_layers == 1:
            return nn.ModuleList([GCNConv(in_channels, out_channels, **common_kwargs)])

        layers = [GCNConv(in_channels, hidden_channels, **common_kwargs)]
        layers.extend(
            GCNConv(hidden_channels, hidden_channels, **common_kwargs)
            for _ in range(num_layers - 2)
        )
        layers.append(GCNConv(hidden_channels, out_channels, **common_kwargs))

        return nn.ModuleList(layers)

class GCNConfig(TypedDict):
    """
    Configuration for the GCN model.

    Attributes:
        in_channels: Dimension of the input node embeddings to the GCN layers.
        out_channels: Dimension of the output node embeddings from the GCN layers.
        hidden_channels: Dimension of the hidden node embeddings in the GCN layers.
        num_layers: Number of GCN layers. Must be at least 1. Defaults to ``2``.
        drop_rate: Dropout rate applied after each GCN layer (except the last one).
            Defaults to ``0.0`` (no dropout).
        activation_fn: Activation function to use after each hidden layer.
            Defaults to ``nn.ReLU``.
        activation_fn_kwargs: Keyword arguments for the activation function.
            Defaults to empty dict.
        bias: Whether to include a bias term in the GCN layers. Defaults to ``True``.
        improved: Whether to use the improved version of GCNConv. Defaults to ``False``.
        add_self_loops: Whether to add self-loops to the input graph. Defaults to ``True``.
        normalize: Whether to symmetrically normalize the adjacency matrix in GCNConv.
            Defaults to ``True``.
        cached: Whether to cache the normalized adjacency matrix in GCNConv.
            Only applicable if the graph structure does not change between epochs.
            Defaults to ``False``.
    """

    in_channels: int
    out_channels: int
    hidden_channels: NotRequired[int]
    num_layers: NotRequired[int]
    drop_rate: NotRequired[float]
    bias: NotRequired[bool]
    activation_fn: NotRequired[ActivationFn]
    activation_fn_kwargs: NotRequired[dict]
    improved: NotRequired[bool]
    add_self_loops: NotRequired[bool]
    normalize: NotRequired[bool]
    cached: NotRequired[bool]

class HGNN(nn.Module):
    """
    HGNN performs spectral convolution directly on the hypergraph structure using the
    node-hyperedge incidence matrix, without any reduction to a pairwise graph.
    Unlike HyperGCN (which approximates each hyperedge by selecting representative pairwise
    edges via random projection), HGNN preserves all higher-order relationships by passing
    messages through the full incidence structure: nodes -> hyperedges -> nodes.

    References:
        - Proposed in [Hypergraph Neural Networks](https://arxiv.org/pdf/1809.09401) (AAAI 2019).
        - Reference implementation: [Code](https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hgnn.html#HGNN).

    Attributes:
        layers: Two stacked ``HGNNConv`` layers.
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
    ):
        """
        Initialize the HGNN model.

        Args:
            in_channels: Number of input channels.
            hidden_channels: Number of hidden channels.
            num_classes: Number of output channels.
            bias: If set to ``False``, the layer will not learn the bias parameter.
                Defaults to ``True``.
            use_batch_normalization: If set to ``True``, layers will use batch normalization.
                Defaults to ``False``.
            drop_rate: Dropout ratio. Defaults to ``0.5``.
        """
        super().__init__()

        self.layers: nn.ModuleList = nn.ModuleList(
            [
                HGNNConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                ),
                HGNNConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply two stacked ``HGNNConv`` layers to produce node embeddings.

        The first layer applies ReLU + dropout and maps ``in_channels -> hidden_channels``.
        The second layer is the output layer (no activation/dropout) and maps
        ``hidden_channels -> num_classes``.

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge incidence in COO format. Size ``(2, num_incidences)``,
                where row 0 contains node IDs and row 1 contains hyperedge IDs.

        Returns:
            x: The output node feature matrix. Size ``(num_nodes, num_classes)``.
        """
        for layer in self.layers:
            x = layer(x, hyperedge_index)
        return x

class HNHN(nn.Module):
    """
    HNHN performs incidence-based hypergraph convolution with explicit hyperedge embeddings between
    the node -> hyperedge -> node propagation steps.

    References:
        - Proposed in [HNHN: Hypergraph Networks with Hyperedge Neurons](https://arxiv.org/abs/2006.12278) paper.
        - Reference implementation: [Code](https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hnhn.html#HNHN).

    Attributes:
        layers: Two stacked ``HNHNConv`` layers.
    """  # noqa: E501

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
    ):
        """
        Initialize the HNHN model.

        Args:
            in_channels: Number of input channels.
            hidden_channels: Number of hidden channels.
            num_classes: Number of output channels.
            bias: If set to ``False``, the layer will not learn the bias parameter.
                Defaults to ``True``.
            use_batch_normalization: If set to ``True``, layers will use batch normalization.
                Defaults to ``False``.
            drop_rate: Dropout ratio. Defaults to ``0.5``.
        """
        super().__init__()

        self.layers: nn.ModuleList = nn.ModuleList(
            [
                HNHNConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                ),
                HNHNConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply two stacked ``HNHNConv`` layers to produce node embeddings.

        Args:
            x: Input node feature matrix of size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge incidence in COO format of size ``(2, num_incidences)``.

        Returns:
            x: The output node feature matrix of size ``(num_nodes, num_classes)``.
        """
        for layer in self.layers:
            x = layer(x, hyperedge_index)
        return x

class HGNNP(nn.Module):
    """
    HGNN+ performs hypergraph convolution with two-stage mean aggregation using the
    incidence structure directly: nodes -> hyperedges -> nodes.

    References:
        - Proposed in [HGNN+: General Hypergraph Neural Networks](https://ieeexplore.ieee.org/document/9795251) paper (IEEE T-PAMI 2022).
        - Reference implementation: [Code](https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hgnnp.html#HGNNP).

    Attributes:
        layers: Two stacked ``HGNNPConv`` layers.
    """  # noqa: E501

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
    ):
        """
        Initialize the HGNN+ model.

        Args:
            in_channels: Number of input channels.
            hidden_channels: Number of hidden channels.
            num_classes: Number of output channels.
            bias: If set to ``False``, the layer will not learn the bias parameter.
                Defaults to ``True``.
            use_batch_normalization: If set to ``True``, layers will use batch normalization.
                Defaults to ``False``.
            drop_rate: Dropout ratio. Defaults to ``0.5``.
        """
        super().__init__()

        self.layers: nn.ModuleList = nn.ModuleList(
            [
                HGNNPConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                ),
                HGNNPConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    is_last=True,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply two stacked ``HGNNPConv`` layers to produce node embeddings.

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: Hyperedge incidence in COO format. Size ``(2, num_incidences)``,
                where row 0 contains node IDs and row 1 contains hyperedge IDs.

        Returns:
            x: The output node feature matrix. Size ``(num_nodes, num_classes)``.
        """
        for layer in self.layers:
            x = layer(x, hyperedge_index)
        return x

class HyperGCN(nn.Module):
    """
    HyperGCN approximates each hyperedge of the hypergraph by a set of pairwise edges connecting the
    vertices of the hyperedge and treats the learning problem as a graph learning problem on the
    approximation.

    References:
        - Proposed in [HyperGCN: A New Method of Training Graph Convolutional Networks on Hypergraphs](https://dl.acm.org/doi/10.5555/3454287.3454422) paper (NeurIPS 2019).
        - Code of the paper: [source](https://github.com/malllabiisc/HyperGCN).
        - Reference implementation: [source](https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/models/hypergraphs/hypergcn.html#HyperGCN).

    Attributes:
        fast: Whether to cache the reduced graph smoothing matrix.
        use_mediator: Whether to use mediator edges during hypergraph reduction.
        cached_gcn_laplacian_matrix: Cached normalized GCN Laplacian matrix.
        seed: Optional random seed for hypergraph reduction.
        layers: Two stacked ``HyperGCNConv`` layers.
    """  # noqa: E501

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        num_classes: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        use_mediator: bool = False,
        fast: bool = True,
        seed: int | None = None,
    ):
        """
        Initialize the HyperGCN model.

        Args:
            in_channels: The number of input channels.
            hidden_channels: The number of hidden channels.
            num_classes: The number of classes of the classification task as HyperGCB is a
                node classification model.
            bias: If set to ``False``, the layer will not learn the bias parameter.
                Defaults to ``True``.
            use_batch_normalization: If set to ``True``, layers will use batch normalization.
                Defaults to ``False``.
            drop_rate: Dropout ratio. Defaults to ``0.5``.
            use_mediator: Whether to use mediator to transform the hyperedges to edges in the graph.
                Defaults to ``False``.
            fast: If set to ``True``, the transformed graph structure will be computed once from
                the input hypergraph and vertex features, and cached for future use.
                Defaults to ``True``.
            seed: Optional random seed for the random reduction of hyperedges to edges.
                Defaults to ``None``.
        """
        super().__init__()
        self.fast: bool = fast
        self.use_mediator: bool = use_mediator
        self.cached_gcn_laplacian_matrix: Tensor | None = None
        self.seed: int | None = seed

        self.layers: nn.ModuleList = nn.ModuleList(
            [
                HyperGCNConv(
                    in_channels=in_channels,
                    out_channels=hidden_channels,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    drop_rate=drop_rate,
                    use_mediator=use_mediator,
                    seed=self.seed,
                ),
                HyperGCNConv(
                    in_channels=hidden_channels,
                    out_channels=num_classes,
                    bias=bias,
                    use_batch_normalization=use_batch_normalization,
                    use_mediator=use_mediator,
                    is_last=True,
                    seed=self.seed,
                ),
            ]
        )

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        The forward function.

        Args:
            x: Input node feature matrix. Size ``(num_nodes, in_channels)``.
            hyperedge_index: The hyperedge indices of the hypergraph. Size ``(2, num_hyperedges)``.

        Returns:
            x: The output node feature matrix. Size ``(num_nodes, num_classes)``.
        """
        if not self.fast:
            for layer in self.layers:
                x = layer(x, hyperedge_index)
            return x

        # If the GCN Laplacian is cached, we need to check if the node feature size has changed
        # with cached_gcn_laplacian_matrix.size(0) != x.size(0), this can happen,
        # for example, due to:
        # adding new negative samples or having validation/test sets with different node features
        should_not_use_cached_gcn_laplacian_matrix = (
            self.cached_gcn_laplacian_matrix is None  # Not cached yet
            or self.cached_gcn_laplacian_matrix.size(0)
            != x.size(0)  # Node feature size has changed
        )

        if should_not_use_cached_gcn_laplacian_matrix:
            edge_index, edge_weights = HyperedgeIndex(
                hyperedge_index
            ).reduce_to_edge_index_on_random_direction(
                x=x,
                with_mediators=self.use_mediator,
                return_weights=True,
                seed=self.seed,
            )

            self.cached_gcn_laplacian_matrix = EdgeIndex(
                edge_index=edge_index,
                edge_weights=edge_weights,
            ).get_sparse_normalized_gcn_laplacian(num_nodes=x.size(0))

        for layer in self.layers:
            x = layer(x, hyperedge_index, gcn_laplacian_matrix=self.cached_gcn_laplacian_matrix)
        return x