modeling/models/acm_gineconv.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 GraphNorm(nn.Module):
    """GraphNorm (Cai et al., ICML 2021): per-graph normalization with a
    learnable mean-shift, applied to node states between GNN layers.

    Faithful to the official implementation (lsj2408/GraphNorm):
        sub = x - alpha * mean_g(x)          # per-graph mean, learnable shift
        out = gamma * sub / sqrt(mean_g(sub^2) + eps) + beta

    where mean_g is computed within each graph of the batch. Unlike BatchNorm
    (cross-graph batch statistics -> noisy on heterogeneous graphs), GraphNorm
    normalizes each graph independently, which is the principled choice for
    batches of heterogeneous cell-graphs.
    """

    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))   # scale (weight)
        self.beta = nn.Parameter(torch.zeros(dim))    # shift (bias)
        self.alpha = nn.Parameter(torch.ones(dim))    # learnable mean-shift (mean_scale)

    def reset_parameters(self):
        nn.init.ones_(self.gamma)
        nn.init.zeros_(self.beta)
        nn.init.ones_(self.alpha)

    def forward(self, x: Tensor, batch: Tensor = None) -> Tensor:
        if batch is None:
            batch = torch.zeros(x.size(0), dtype=torch.long, device=x.device)
        num_graphs = int(batch.max().item()) + 1
        # Per-graph mean, broadcast back to nodes.
        mean = scatter(x, batch, 0, num_graphs, reduce="mean")[batch]
        sub = x - self.alpha * mean
        var = scatter(sub * sub, batch, 0, num_graphs, reduce="mean")[batch]
        std = (var + self.eps).sqrt()
        return self.gamma * sub / std + self.beta

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(dim={self.dim})"


class ACM_GINEConv(MessagePassing):
    """Single ACM-GINEConv layer: GINEConv-style message passing + ACM filtering.

    Messages incorporate static edge features via addition before ReLU
    (following the GINEConv formulation), then aggregated messages are
    decomposed into low-pass, high-pass, and identity channels with
    learned adaptive mixing.

    Message:  m_ij = a_ij * ReLU(H_j + E_ji)
    where E_ji are static encoded edge features (no per-layer edge MLP).
    """

    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-GINEConv 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] (used for
                degree normalization).
            edge_feat: Static edge features [E, hidden_dim] encoded once
                at the start; same features reused at every layer.
            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 normalization using sum of edge weights
        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_GINEConv"
        )

        # ACM frequency decomposition
        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))

        # Adaptive mixing
        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:
        """GINEConv-style message: m_ij = a_ij * ReLU(H_j + E_ji)."""
        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_GINEConv_model(nn.Module):
    """Multi-layer ACM-GINEConv model with static edge feature integration.

    Unlike ACM_GIN_model (which learns per-layer edge MLPs), this model
    encodes edge features once via a linear projection and reuses them
    at every GNN layer. This follows the GINEConv philosophy:
      - Edge features are incorporated into messages via ReLU(H_j + E_ji)
      - No per-layer edge update MLP (significant memory savings)
      - ACM filtering and adaptive mixing remain unchanged

    This is a simpler, more memory-efficient alternative suitable for
    pretraining on a single GPU where the full edge-update model may OOM.
    """

    def __init__(
        self,
        in_dim,
        out_dim,
        num_layers,
        hidden_dim,
        edge_in_dim,
        batchnorm,
        activation="relu",
        norm_type="none",
        input_norm="none",
        edge_distance_in_proj=True,
    ):
        super(ACM_GINEConv_model, self).__init__()
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.gnn_batchnorm = batchnorm
        if norm_type not in ("none", "layer", "graph"):
            raise ValueError(f"Unknown norm_type {norm_type!r}; expected none/layer/graph")
        self.norm_type = norm_type
        if input_norm not in ("none", "batch", "layer"):
            raise ValueError(f"Unknown input_norm {input_norm!r}; expected none/batch/layer")
        self.input_norm_type = input_norm
        self.out_dim = out_dim

        # Column 0 of edge_attr is the spatial distance. When False, exclude it
        # from the edge feature projection: the distance is then used ONLY as the
        # per-edge aggregation weight a_ij (and degree norm), not mixed into the
        # additive edge feature E_ji. This also removes the raw-distance scale
        # from the projection input.
        self.edge_distance_in_proj = edge_distance_in_proj
        edge_proj_in = edge_in_dim if edge_distance_in_proj else edge_in_dim - 1

        # Project raw node and edge features into hidden_dim once.
        # Edge features remain static after this projection.
        self.node_input_proj = nn.Linear(in_dim, hidden_dim)
        self.edge_input_proj = nn.Linear(edge_proj_in, hidden_dim)

        # Optional LEARNABLE normalization of the projected input embedding
        # (applied right after node_input_proj). Makes the input conditioning
        # adaptive; the static json NormalizeData still defines the fixed SCE
        # target. "batch" = online per-feature stats; "layer" = per-node.
        if input_norm == "batch":
            self.input_norm = nn.BatchNorm1d(hidden_dim)
        elif input_norm == "layer":
            self.input_norm = nn.LayerNorm(hidden_dim)
        else:
            self.input_norm = None

        # Optional per-layer normalization on node states BETWEEN conv layers
        # (the research-conventional spot for LayerNorm/GraphNorm). Anti-collapse
        # measure: stops variance concentrating onto a single PCA direction.
        #   "layer" -> nn.LayerNorm  (per-node, across features)
        #   "graph" -> GraphNorm     (per-graph, across nodes; needs batch)
        # (BatchNorm lives inside the channel MLPs via `batchnorm`, the GIN spot.)
        self.layer_norms = nn.ModuleList() if norm_type in ("layer", "graph") else None

        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.activation_name = activation

        for i in range(self.num_layers):
            # --- 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))

            # --- ACM channel MLPs ---
            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_GINEConv(
                    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],
                )
            )

            if self.norm_type == "layer":
                self.layer_norms.append(nn.LayerNorm(local_out_dim))
            elif self.norm_type == "graph":
                self.layer_norms.append(GraphNorm(local_out_dim))

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

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

        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.
            batch: Node->graph assignment [N]; required only when
                norm_type == "graph" (GraphNorm normalizes per graph).
            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
        edge_weight = edge_attr[:, 0]

        # Project node and edge features into hidden_dim ONCE
        x = self.node_input_proj(x)
        if self.input_norm is not None:
            x = self.input_norm(x)  # learnable input-embedding normalization
        # Distance (col 0) stays as edge_weight only when excluded from the proj.
        edge_proj_input = edge_attr if self.edge_distance_in_proj else edge_attr[:, 1:]
        edge_feat = self.edge_input_proj(edge_proj_input)  # Static: reused every layer

        outs = []
        for i in range(self.num_layers):
            # GINEConv-style message passing with ACM filtering
            # Edge features are static — no per-layer edge update
            x = self.ACM_convs[i](
                x=x,
                edge_index=edge_index,
                edge_weight=edge_weight,
                edge_feat=edge_feat,
            )
            if self.layer_norms is not None:
                if self.norm_type == "graph":
                    x = self.layer_norms[i](x, batch)
                else:
                    x = self.layer_norms[i](x)
            outs.append(x)

        if return_hidden:
            return x, outs
        else:
            return x


if __name__ == "__main__":
    model = ACM_GINEConv_model(46, 46, 2, 256, 74, True)
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"ACM_GINEConv_model parameters: {num_params:,}")