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:,}")