modeling/models/acm_gin.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from torch import Tensor
from typing import Union
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn.inits import reset
from torch_geometric.typing import OptPairTensor, Size
from torch_geometric.utils import scatter

from .utils import create_activation


class ACM_GIN(MessagePassing):
    """Single ACM-GIN convolution layer with edge-aware message passing.

    The message from node j to node i incorporates the intermediate edge
    feature e_ij_prime (precomputed by the outer model) alongside the
    scalar spatial weight a_ij used for degree normalization:

        m_ij = a_ij * ReLU(H_j + e_ij_prime)
    """

    def __init__(
        self,
        nn_lowpass: torch.nn.Module,
        nn_highpass: torch.nn.Module,
        nn_fullpass: torch.nn.Module,
        nn_lowpass_proj: torch.nn.Module,
        nn_highpass_proj: torch.nn.Module,
        nn_fullpass_proj: torch.nn.Module,
        nn_mix: torch.nn.Module,
        T: float = 3.0,
        **kwargs,
    ):
        kwargs.setdefault("aggr", "add")
        super().__init__(**kwargs)
        self.nn_lowpass = nn_lowpass
        self.nn_highpass = nn_highpass
        self.nn_fullpass = nn_fullpass
        self.nn_lowpass_proj = nn_lowpass_proj
        self.nn_highpass_proj = nn_highpass_proj
        self.nn_fullpass_proj = nn_fullpass_proj
        self.nn_mix = nn_mix
        self.sigmoid = torch.nn.Sigmoid()
        self.softmax = torch.nn.Softmax(dim=1)
        self.T = T
        self.reset_parameters()

    def reset_parameters(self):
        reset(self.nn_lowpass)
        reset(self.nn_highpass)
        reset(self.nn_fullpass)
        reset(self.nn_lowpass_proj)
        reset(self.nn_highpass_proj)
        reset(self.nn_fullpass_proj)
        reset(self.nn_mix)

    def forward(
        self,
        x: Union[Tensor, OptPairTensor],
        edge_index: Tensor,
        edge_weight: Tensor,
        edge_feat: Tensor,
        size: Size = None,
    ) -> Tensor:
        """Forward pass of a single ACM-GIN layer.

        Args:
            x: Node features [N, hidden_dim] or (x_src, x_dst) pair.
            edge_index: Edge indices [2, E].
            edge_weight: Scalar spatial distance per edge [E] (first column
                of the original edge_attr, used for degree normalization).
            edge_feat: Intermediate edge features [E, hidden_dim], i.e.
                e_ij_prime precomputed by the edge MLP in the outer model.
            size: Optional bipartite graph size.
        """
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)

        # propagate_type: (x: OptPairTensor, edge_weight: Tensor, edge_feat: Tensor)
        out = self.propagate(
            edge_index, x=x, edge_weight=edge_weight, edge_feat=edge_feat, size=size
        )

        # Degree here is the sum of edge weights, not the neighbour count as in
        # standard GIN. Nodes whose weights sum to zero are handled below.
        deg = scatter(edge_weight, edge_index[1], 0, out.size(0), reduce="sum")
        deg_inv = 1.0 / deg
        deg_inv.masked_fill_(deg_inv == float("inf"), 0)
        out = deg_inv.view(-1, 1) * out

        x_r = x[1]
        assert x_r is not None, (
            "Target node features (x_r) must not be None for ACM_GIN"
        )
        out_lowpass = (x_r + out) / 2.0
        out_highpass = (x_r - out) / 2.0

        # compute embeddings for each filter
        out_lowpass = self.nn_lowpass(out_lowpass)
        out_highpass = self.nn_highpass(out_highpass)
        out_fullpass = self.nn_fullpass(x_r)
        # compute importance weights per filter
        alpha_lowpass = self.sigmoid(self.nn_lowpass_proj(out_lowpass))
        alpha_highpass = self.sigmoid(self.nn_highpass_proj(out_highpass))
        alpha_fullpass = self.sigmoid(self.nn_fullpass_proj(out_fullpass))
        alpha_cat = torch.concat([alpha_lowpass, alpha_highpass, alpha_fullpass], dim=1)
        alpha_cat = self.softmax(self.nn_mix(alpha_cat / self.T))

        out = alpha_cat[:, 0].view(-1, 1) * out_lowpass
        out = out + alpha_cat[:, 1].view(-1, 1) * out_highpass
        out = out + alpha_cat[:, 2].view(-1, 1) * out_fullpass

        return out

    def message(self, x_j: Tensor, edge_weight: Tensor, edge_feat: Tensor) -> Tensor:
        """Edge-aware message: m_ij = a_ij * ReLU(H_j + e_ij_prime)."""
        return edge_weight.view(-1, 1) * F.relu(x_j + edge_feat)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}("
            f"nn_lowpass={self.nn_lowpass}, "
            f"nn_highpass={self.nn_highpass}, "
            f"nn_fullpass={self.nn_fullpass})"
        )


class ACM_GIN_model(nn.Module):
    """Multi-layer ACM-GIN model with edge-aware message passing.

    Both node and edge features are projected into hidden_dim at the start
    (via ``node_input_proj`` and ``edge_input_proj``).  This ensures
    uniform dimensions throughout, so every edge MLP receives 3 * hidden_dim
    and the edge residual connection is valid from layer 0 onward.

    At each layer k the model:
      1. Computes intermediate edge features via an edge MLP:
         e_ij_prime = MLP_edge(H_i || H_j || E_ij)
      2. Updates edge state with a residual connection:
         E_ij^(k) = E_ij^(k-1) + e_ij_prime
      3. Passes messages using the scalar spatial weight and the
         intermediate edge features:
         m_ij = a_ij * ReLU(H_j + e_ij_prime)
      4. Applies ACM channel mixing (low/high/full-pass) on the
         aggregated messages.
    """

    def __init__(
        self,
        in_dim,
        out_dim,
        num_layers,
        hidden_dim,
        edge_in_dim,
        batchnorm,
        activation="relu",
    ):
        super(ACM_GIN_model, self).__init__()
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.gnn_batchnorm = batchnorm
        self.out_dim = out_dim

        # Project raw node and edge features into hidden_dim so that both
        # live in the same space from the very start.  This enables the
        # edge residual connection at every layer (including layer 0) and
        # keeps all edge MLP input dimensions uniform at 3 * hidden_dim.
        self.node_input_proj = nn.Linear(in_dim, hidden_dim)
        self.edge_input_proj = nn.Linear(edge_in_dim, hidden_dim)

        self.ACM_convs = nn.ModuleList()
        self.nns_lowpass = nn.ModuleList()
        self.nns_highpass = nn.ModuleList()
        self.nns_fullpass = nn.ModuleList()
        self.nns_lowpass_proj = nn.ModuleList()
        self.nns_highpass_proj = nn.ModuleList()
        self.nns_fullpass_proj = nn.ModuleList()
        self.nns_mix = nn.ModuleList()
        self.edge_mlps = nn.ModuleList()

        self.activation_name = activation

        for i in range(self.num_layers):
            # --- Edge MLP for this layer ---
            # Both nodes and edges have been projected to hidden_dim before
            # the loop, so the input is always (src + dst + edge_state) =
            # 3 * hidden_dim for every layer.
            edge_mlp_in = 3 * hidden_dim

            if self.gnn_batchnorm:
                self.edge_mlps.append(
                    nn.Sequential(
                        nn.Linear(edge_mlp_in, hidden_dim),
                        nn.BatchNorm1d(hidden_dim),
                        create_activation(activation),
                        nn.Linear(hidden_dim, hidden_dim),
                        nn.BatchNorm1d(hidden_dim),
                        create_activation(activation),
                    )
                )
            else:
                self.edge_mlps.append(
                    nn.Sequential(
                        nn.Linear(edge_mlp_in, hidden_dim),
                        create_activation(activation),
                        nn.Linear(hidden_dim, hidden_dim),
                        create_activation(activation),
                    )
                )

            # --- Projection modules to compute importance weights ---
            for channel_proj_module in [
                self.nns_lowpass_proj,
                self.nns_highpass_proj,
                self.nns_fullpass_proj,
            ]:
                if i == self.num_layers - 1:
                    channel_proj_module.append(nn.Linear(self.out_dim, 1))
                else:
                    channel_proj_module.append(nn.Linear(self.hidden_dim, 1))

            # --- Weights mixing module as attention mechanism ---
            self.nns_mix.append(nn.Linear(3, 3))

            # --- GIN channel MLPs ---
            # After node_input_proj, all nodes are hidden_dim, so
            # local_input_dim is always hidden_dim.
            local_input_dim = self.hidden_dim

            if i == self.num_layers - 1:
                local_out_dim = self.out_dim
            else:
                local_out_dim = self.hidden_dim

            for channel_module in [
                self.nns_lowpass,
                self.nns_highpass,
                self.nns_fullpass,
            ]:
                if self.gnn_batchnorm:
                    sequential = nn.Sequential(
                        nn.Linear(local_input_dim, self.hidden_dim),
                        nn.BatchNorm1d(self.hidden_dim),
                        create_activation(self.activation_name),
                        nn.Linear(self.hidden_dim, local_out_dim),
                        nn.BatchNorm1d(local_out_dim),
                        create_activation(self.activation_name),
                    )
                else:
                    sequential = nn.Sequential(
                        nn.Linear(local_input_dim, self.hidden_dim),
                        create_activation(self.activation_name),
                        nn.Linear(self.hidden_dim, local_out_dim),
                        create_activation(self.activation_name),
                    )

                channel_module.append(sequential)

            self.ACM_convs.append(
                ACM_GIN(
                    nn_lowpass=self.nns_lowpass[i],
                    nn_highpass=self.nns_highpass[i],
                    nn_fullpass=self.nns_fullpass[i],
                    nn_lowpass_proj=self.nns_lowpass_proj[i],
                    nn_highpass_proj=self.nns_highpass_proj[i],
                    nn_fullpass_proj=self.nns_fullpass_proj[i],
                    nn_mix=self.nns_mix[i],
                )
            )

    def reset_parameters(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                m.reset_parameters()
            elif isinstance(m, nn.BatchNorm1d):
                m.reset_parameters()

    def forward(self, x, edge_index, edge_attr, batch=None, return_hidden=False):
        """Forward pass through all ACM-GIN layers with edge updates.

        `batch` is accepted for API parity with ACM_GINEConv_model (GraphNorm);
        it is unused here.

        Args:
            x: Node features [N, in_dim].
            edge_index: Edge indices [2, E].
            edge_attr: Edge features [E, edge_in_dim]. The first column
                (index 0) is the scalar spatial distance used for degree
                normalization; the full vector evolves through layers via
                the edge MLPs.
            return_hidden: If True, also return all intermediate node states.

        Returns:
            x: Final node embeddings [N, out_dim].
            outs: (optional) List of node states after each layer.
        """
        # Extract scalar spatial distance for degree normalization
        # (stays fixed across layers)
        edge_weight = edge_attr[:, 0]

        # Project node and edge features from raw dims to hidden_dim
        x = self.node_input_proj(x)
        edge_state = self.edge_input_proj(edge_attr)

        outs = []
        for i in range(self.num_layers):
            # Step 1: Compute intermediate edge features
            src, dst = edge_index
            edge_mlp_input = torch.cat([x[src], x[dst], edge_state], dim=-1)
            e_ij_prime = self.edge_mlps[i](edge_mlp_input)

            # Step 2: Update edge state with residual (safe at every layer
            # because both edge_state and e_ij_prime are hidden_dim)
            edge_state = edge_state + e_ij_prime

            # Step 3-4: Edge-aware ACM message passing
            x = self.ACM_convs[i](
                x=x,
                edge_index=edge_index,
                edge_weight=edge_weight,
                edge_feat=e_ij_prime,
            )
            outs.append(x)

        if return_hidden:
            return x, outs
        else:
            return x


if __name__ == "__main__":
    acm_gin = ACM_GIN_model(46, 46, 2, 256, 74, True)
    print(sum(p.numel() for p in acm_gin.parameters() if p.requires_grad))