NN¶

class HyperedgeAggregator:
    """
    Pool node embeddings into hyperedge embeddings using the incidence structure.

    Each node-hyperedge incidence selects one node embedding row, then reduces
    those rows per hyperedge with the requested scatter aggregation.

    Attributes:
        hyperedge_index_wrapper: Wrapper around the hyperedge incidence tensor.
        node_embeddings: Node embedding matrix of size ``(num_nodes, num_channels)``.
        num_hyperedges: Optional explicit hyperedge count.
            When provided, the pooled output preserves empty hyperedges that
            do not appear in ``hyperedge_index``.
    """

    def __init__(
        self,
        hyperedge_index: Tensor,
        node_embeddings: Tensor,
        num_hyperedges: int | None = None,
    ):
        """
        Initialize the hyperedge aggregator.

        Args:
            hyperedge_index: Hyperedge incidence in COO format of size ``(2, num_incidences)``.
            node_embeddings: Node embedding matrix of size ``(num_nodes, num_channels)``.
            num_hyperedges: Optional explicit hyperedge count. Defaults to ``None``.
        """
        self.hyperedge_index_wrapper: HyperedgeIndex = HyperedgeIndex(hyperedge_index)
        self.node_embeddings: Tensor = node_embeddings
        self.num_hyperedges: int | None = num_hyperedges

    def pool(self, aggregation: Literal["maxmin", "max", "min", "mean", "mul", "sum"]) -> Tensor:
        """
        Aggregate node embeddings for each hyperedge.

        ``hyperedge_index`` is the COO encoding of the nonzero entries of ``H``,
        so ``hyperedge_index[0, k] = v`` and ``hyperedge_index[1, k] = e`` means ``H[v, e] = 1``
        for incidence ``k``.

        Let ``H`` be the binary incidence matrix of shape ``(num_nodes, num_hyperedges)``
        and let ``X`` be the node embedding matrix of shape ``(num_nodes, num_channels)``.
        This method pools node features into hyperedge features using the incidence pattern in
        ``H``.

        Aggregations:
            - ``aggregation="sum"`` computes the equivalent of the standard
                sparse matrix product ``H^T X``.
            - ``aggregation="mean"`` computes ``D_e^{-1} H^T X``, where
                ``D_e[e, e] = sum_v H[v, e]`` is the hyperedge cardinality matrix.
            - ``aggregation in {"max", "min", "mul"}`` uses the same sparsity pattern as ``H^T X``,
                but replaces the summation over incident nodes with a channel-wise ``max``, ``min``,
                or product reduction.
            - ``aggregation="maxmin"`` computes the channel-wise range ``max - min``
                for each hyperedge.

        Examples:
            >>> hyperedge_index = [[0, 1, 2, 2, 3],
            ...                    [0, 0, 0, 1, 1]]
            >>> node_embeddings = [[1, 10], [2, 20], [3, 30], [4, 40]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("mean")
            ... [[2, 20], [3.5, 35]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("sum")
            ... [[6, 60], [7, 70]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("max")
            ... [[3, 30], [4, 40]]
            >>> HyperedgeAggregator(hyperedge_index, node_embeddings).pool("maxmin")
            ... [[2, 20], [1, 10]]

        Args:
            aggregation: Reduction applied across the nodes belonging to each hyperedge.

        Returns:
            hyperedge_embeddings: A hyperedge embedding matrix of
                shape ``(num_hyperedges, num_channels)``.
        """
        # Gather the embeddings for each incidence.
        # A node appearing in multiple hyperedges is repeated, once per incidence.
        # Example: node_embeddings = [[1, 10],  # node 0
        #                             [2, 20],  # node 1
        #                             [3, 30],  # node 2
        #                             [4, 40]]  # node 3
        #          -> all_node_ids = [0, 1, 2, 2, 3]
        #          -> incidence_node_embeddings = [[1, 10],  # node 0 for hyperedge 0
        #                                          [2, 20],  # node 1 for hyperedge 0
        #                                          [3, 30],  # node 2 for hyperedge 0
        #                                          [3, 30],  # node 2 for hyperedge 1
        #                                          [4, 40]]  # node 3 for hyperedge 1
        #             shape: (num_incidences, num_channels)
        incidence_node_embeddings = self.node_embeddings[self.hyperedge_index_wrapper.all_node_ids]

        # Scatter-aggregate node embeddings into hyperedge embeddings.
        # Example: with aggregation="sum":
        #          [[1+2+3, 10+20+30],  # hyperedge 0 contains node 0, 1, 2
        #          [3+4, 30+40]]        # hyperedge 1 contains node 2, 3
        #          shape: (num_hyperedges, num_channels)
        #          with aggregation="max":
        #          [[max(1, 2, 3), max(10, 20, 30)],  # hyperedge 0 contains node 0, 1, 2
        #           [max(3, 4), max(30, 40)]]         # hyperedge 1 contains node 2, 3
        #          shape: (num_hyperedges, num_channels)
        num_hyperedges = (
            self.num_hyperedges
            if self.num_hyperedges is not None
            else self.hyperedge_index_wrapper.num_hyperedges
        )

        if aggregation == "maxmin":
            return maxmin_scatter(
                src=incidence_node_embeddings,
                index=self.hyperedge_index_wrapper.all_hyperedge_ids,
                dim=0,  # scatter along the hyperedge dimension
                dim_size=num_hyperedges,
            )

        return scatter(
            src=incidence_node_embeddings,
            index=self.hyperedge_index_wrapper.all_hyperedge_ids,
            dim=0,  # scatter along the hyperedge dimension
            dim_size=num_hyperedges,
            reduce=aggregation,
        )

class NodeAggregator:
    """
    Pool hyperedge embeddings into node embeddings using the incidence structure.

    Each node-hyperedge incidence selects one hyperedge embedding row, then
    reduces those rows per node with the requested scatter aggregation.

    Attributes:
        hyperedge_index_wrapper: Wrapper around the hyperedge incidence tensor.
        hyperedge_embeddings: Hyperedge embedding matrix of size ``(num_hyperedges, num_channels)``.
        num_nodes: Optional explicit node count. When provided, the pooled output preserves
            isolated nodes that do not appear in ``hyperedge_index``. Defaults to ``None``.
    """

    def __init__(
        self,
        hyperedge_index: Tensor,
        hyperedge_embeddings: Tensor,
        num_nodes: int | None = None,
    ):
        """
        Initialize the node aggregator.

        Args:
            hyperedge_index: Hyperedge incidence in COO format of size ``(2, num_incidences)``.
            hyperedge_embeddings: Hyperedge embedding matrix of size
                ``(num_hyperedges, num_channels)``.
            num_nodes: Optional explicit node count. Defaults to ``None``.
        """
        self.hyperedge_index_wrapper: HyperedgeIndex = HyperedgeIndex(hyperedge_index)
        self.hyperedge_embeddings: Tensor = hyperedge_embeddings
        self.num_nodes: int | None = num_nodes

    def pool(self, aggregation: Literal["maxmin", "max", "min", "mean", "mul", "sum"]) -> Tensor:
        """
        Aggregate hyperedge embeddings for each node.

        ``hyperedge_index`` is the COO encoding of the nonzero entries of ``H``,
        so ``hyperedge_index[0, k] = v`` and ``hyperedge_index[1, k] = e`` means ``H[v, e] = 1``
        for incidence ``k``.

        Let ``H`` be the incidence matrix of shape ``(num_nodes, num_hyperedges)``
        and let ``E`` be the hyperedge embedding matrix of shape ``(num_hyperedges, num_channels)``.
        This method pools hyperedge features into node features using the incidence pattern
        in ``H``.

        Aggregations:
            - ``aggregation="sum"`` computes the equivalent of the standard
                sparse matrix product ``H E``.
            - ``aggregation="mean"`` computes ``D_v^{-1} H E``, where ``D_v[v, v] = sum_e H[v, e]``
                is the node degree matrix.
            - ``aggregation in {"max", "min", "mul"}`` uses the same sparsity pattern as ``H E``,
                but replaces the summation over incident hyperedges with a channel-wise
                ``max``, ``min``, or product reduction.

        Examples:
            >>> hyperedge_index = [[0, 1, 1, 2],
            ...                    [0, 0, 1, 1]]
            >>> hyperedge_embeddings = [[10, 100], [20, 200]]
            >>> NodeAggregator(hyperedge_index, hyperedge_embeddings).pool("mean")
            ... [[10, 100], [15, 150], [20, 200]]
            >>> NodeAggregator(hyperedge_index, hyperedge_embeddings).pool("sum")
            ... [[10, 100], [30, 300], [20, 200]]
            >>> NodeAggregator(hyperedge_index, hyperedge_embeddings).pool("max")
            ... [[10, 100], [20, 200], [20, 200]]

        Args:
            aggregation: Reduction applied across the hyperedges incident to each node.

        Returns:
            node_embeddings: A node embedding matrix of shape ``(num_nodes, num_channels)``.
        """
        # Gather the embeddings for each incidence.
        # A hyperedge appearing in multiple node incidences is repeated, once per incidence.
        # Example: hyperedge_embeddings = [[10, 100],  # hyperedge 0
        #                                  [20, 200]]  # hyperedge 1
        #          -> all_hyperedge_ids = [0, 0, 1, 1]
        #          -> incidence_hyperedge_embeddings = [[10, 100],   # hyperedge 0 for node 0
        #                                               [10, 100],   # hyperedge 0 for node 1
        #                                               [20, 200],   # hyperedge 1 for node 1
        #                                               [20, 200]]   # hyperedge 1 for node 2
        #             shape: (num_incidences, num_channels)
        incidence_hyperedge_embeddings = self.hyperedge_embeddings[
            self.hyperedge_index_wrapper.all_hyperedge_ids
        ]
        num_nodes = (
            self.num_nodes if self.num_nodes is not None else self.hyperedge_index_wrapper.num_nodes
        )

        if aggregation == "maxmin":
            return maxmin_scatter(
                src=incidence_hyperedge_embeddings,
                index=self.hyperedge_index_wrapper.all_node_ids,
                dim=0,  # scatter along the node dimension
                dim_size=num_nodes,
            )

        # Scatter-aggregate hyperedge embeddings into node embeddings.
        # Example: with aggregation="sum":
        #          [[10, 100],         # node 0 belongs to hyperedge 0
        #           [10+20, 100+200],  # node 1 belongs to hyperedge 0 and 1
        #           [20, 200]]         # node 2 belongs to hyperedge 1
        #          shape: (num_nodes, num_channels)
        #          with aggregation="max":
        #          [[10, 100],                     # node 0 belongs to hyperedge 0
        #           [max(10, 20), max(100, 200)],  # node 1 belongs to hyperedge 0 and 1
        #           [20, 200]]                     # node 2 belongs to hyperedge 1
        #         shape: (num_nodes, num_channels)
        return scatter(
            src=incidence_hyperedge_embeddings,
            index=self.hyperedge_index_wrapper.all_node_ids,
            dim=0,  # scatter along the node dimension
            dim_size=num_nodes,
            reduce=aggregation,
        )

class HGNNConv(nn.Module):
    """
    References:
        - The HGNNConv layer proposed in [Hypergraph Neural Networks](https://arxiv.org/pdf/1809.09401) paper (AAAI 2019).
        - Reference implementation: [Code](https://deephypergraph.readthedocs.io/en/latest/_modules/dhg/nn/convs/hypergraphs/hgnn_conv.html#HGNNConv).

    Each layer performs: ``X' = sigma(L_HGNN X Theta)``
    where ``L_HGNN = D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2}``
    is the hypergraph Laplacian computed from the incidence matrix H.
    This smooths node features through the hypergraph structure (nodes -> hyperedges -> nodes)
    without reducing to a pairwise graph.

    Unlike ``HyperGCNConv``, which uses a GCN Laplacian on a graph reduced from the hypergraph,
    ``HGNNConv`` operates entirely in hypergraph space and preserves all higher-order relationships.

    Attributes:
        is_last: Whether to skip the final activation and dropout.
            If ``True``, the layer will not apply the final activation and
            dropout functions. Defaults to ``False``.
        batch_norm_1d: Optional batch normalization layer.
        activation_fn: Activation function applied to hidden outputs.
        dropout: Dropout layer applied to hidden outputs. Defaults to ``0.5``.
        theta: Learnable feature projection.
    """  # noqa: E501

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        is_last: bool = False,
    ):
        """
        Initialize the HGNN convolution layer.

        Args:
            in_channels: The number of input channels.
            out_channels: The 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``, the layer will use batch normalization.
                Defaults to ``False``.
            drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
            is_last: If set to ``True``, the layer will not apply the final activation and
                dropout functions. Defaults to ``False``.
        """
        super().__init__()
        self.is_last: bool = is_last
        self.batch_norm_1d: nn.BatchNorm1d | None = (
            nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        )
        self.activation_fn: nn.ReLU = nn.ReLU(inplace=True)
        self.dropout: nn.Dropout = nn.Dropout(drop_rate)
        self.theta: nn.Linear = nn.Linear(in_channels, out_channels, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply one HGNN convolution layer: project features, smooth via hypergraph Laplacian,
        then apply activation, batch norm, and dropout (unless this is the last layer).

        The full per-layer formula is:
            ``X' = sigma( D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2} (X Theta) )``

        where the Laplacian ``L = D_n^{-1/2} H D_e^{-1} H^T D_n^{-1/2}`` is computed from
        the hyperedge_index and can be passed in precomputed as ``hgnn_laplacian_matrix``
        for efficiency when the hypergraph structure does not change across forward passes.

        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, out_channels)``.
        """
        x = self.theta(x)

        smoothing_matrix = HyperedgeIndex(hyperedge_index).get_sparse_hgnn_smoothing_matrix(
            num_nodes=x.size(0),
        )
        x = Hypergraph.smoothing_with_matrix(x, smoothing_matrix)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x

class HGNNPConv(nn.Module):
    """
    References:
        - The HGNNPConv layer 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/nn/convs/hypergraphs/hgnnp_conv.html#HGNNPConv).

    Each layer performs: ``X' = sigma(M_HGNN+ X Theta)`` where
    ``M_HGNN+ = D_v^{-1} H D_e^{-1} H^T`` is the HGNN+ smoothing matrix.

    Unlike ``HGNNConv``, which uses symmetric ``D_v^{-1/2}`` normalization for a
    spectral Laplacian, ``HGNNPConv`` uses plain inverse degrees and performs
    two-stage mean aggregation: nodes -> hyperedges -> nodes.

    Attributes:
        is_last: Whether to skip the final activation and dropout.
            If ``True``, the layer will not apply the final activation and
            dropout functions. Defaults to ``False``.
        batch_norm_1d: Optional batch normalization layer.
        activation_fn: Activation function applied to hidden outputs.
        dropout: Dropout layer applied to hidden outputs. Defaults to ``0.5``.
        theta: Learnable feature projection.
    """  # noqa: E501

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        is_last: bool = False,
    ):
        """
        Initialize the HGNN+ convolution layer.

        Args:
            in_channels: The number of input channels.
            out_channels: The 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``, the layer will use batch normalization.
                Defaults to ``False``.
            drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
            is_last: If set to ``True``, the layer will not apply the final activation and dropout
                functions. Defaults to ``False``.
        """
        super().__init__()
        self.is_last: bool = is_last
        self.batch_norm_1d: nn.BatchNorm1d | None = (
            nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        )
        self.activation_fn: nn.ReLU = nn.ReLU(inplace=True)
        self.dropout: nn.Dropout = nn.Dropout(drop_rate)
        self.theta: nn.Linear = nn.Linear(in_channels, out_channels, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply one HGNN+ convolution layer using row-stochastic hypergraph smoothing.

        The full per-layer formula is:
            ``X' = sigma( D_v^{-1} H D_e^{-1} H^T (X Theta) )``

        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, out_channels)``.
        """
        x = self.theta(x)

        smoothing_matrix = HyperedgeIndex(hyperedge_index).get_sparse_hgnnp_smoothing_matrix(
            num_nodes=x.size(0),
        )
        x = Hypergraph.smoothing_with_matrix(x, smoothing_matrix)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x

class HNHNConv(nn.Module):
    """
    References:
        - The HNHNConv layer 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/nn/convs/hypergraphs/hnhn_conv.html#HNHNConv).

    Attributes:
        is_last: Whether to skip the final activation and dropout.
            If ``True``, the layer will not apply the final activation and
            dropout functions. Defaults to ``False``.
        batch_norm_1d: Optional batch normalization layer.
        activation_fn: Activation function applied to hidden outputs.
        dropout: Dropout layer applied to hidden outputs. Defaults to ``0.5``.
        theta_v2e: Learnable node-to-hyperedge projection.
        theta_e2v: Learnable hyperedge-to-node projection.
    """  # noqa: E501

    __AGGREGATION: Literal["mean"] = "mean"

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        bias: bool = True,
        use_batch_normalization: bool = False,
        drop_rate: float = 0.5,
        is_last: bool = False,
    ):
        """
        Initialize the HNHN convolution layer.

        Args:
            in_channels: The number of input channels.
            out_channels: The 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``, the layer will use batch normalization.
                Defaults to ``False``.
            drop_rate: If set to a positive number, the layer will use dropout. Defaults to ``0.5``.
            is_last: If set to ``True``, the layer will not apply the final activation and
                dropout functions. Defaults to ``False``.
        """
        super().__init__()
        self.is_last: bool = is_last
        self.batch_norm_1d: nn.BatchNorm1d | None = (
            nn.BatchNorm1d(out_channels) if use_batch_normalization else None
        )
        self.activation_fn: nn.ReLU = nn.ReLU(inplace=True)
        self.dropout: nn.Dropout = nn.Dropout(drop_rate)
        self.theta_v2e: nn.Linear = nn.Linear(in_channels, out_channels, bias=bias)
        self.theta_e2v: nn.Linear = nn.Linear(out_channels, out_channels, bias=bias)

    def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
        """
        Apply one HNHN convolution layer using two learned projections around node-to-hyperedge and
        hyperedge-to-node mean aggregation.

        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, out_channels)``.
        """
        x = self.theta_v2e(x)

        hyperedge_embeddings = HyperedgeAggregator(hyperedge_index, x).pool(self.__AGGREGATION)
        hyperedge_embeddings = self.activation_fn(hyperedge_embeddings)
        hyperedge_embeddings = self.theta_e2v(hyperedge_embeddings)

        x = NodeAggregator(
            hyperedge_index=hyperedge_index,
            hyperedge_embeddings=hyperedge_embeddings,
            num_nodes=x.size(0),
        ).pool(self.__AGGREGATION)

        if not self.is_last:
            x = self.activation_fn(x)
            if self.batch_norm_1d is not None:
                x = self.batch_norm_1d(x)
            x = self.dropout(x)

        return x